2D prefix sums and 2D Fenwick trees in Python

Why this matters

Many programs store data in a grid, which is a table with rows and columns. A seating chart, a map of city blocks, a spreadsheet of sales by week and product, and a picture made of pixels are all grids. A common question about a grid is: what is the total inside a rectangle? For example, how many seats are taken in the block from row 2 to row 3 and column 3 to column 4? Or how many sales were made in a range of weeks for a range of products?

A rectangle answer that is computed by looping over every cell is slow when the grid is large and many questions are asked. This lesson shows two ways to make it fast. The first is a two-dimensional prefix table. It answers any rectangle in a fixed amount of work, but it must be rebuilt after every change. The second is a two-dimensional Fenwick tree. It answers rectangles in a small amount of work, and it can update one cell in the same small amount of work. Both ideas extend the one-dimensional versions, which the lesson on counting inversions with a Fenwick tree covers earlier in this module.

Two terms are used throughout. A cell is one box in the grid, identified by its row and column. A rectangle is all cells from row r1 to row r2 and from column c1 to column c2, with both ends included. The code uses 1-based indices, which means the first row is row 1 and the first column is column 1. Python lists start at 0, so the lists in the code have one extra row and one extra column, and the row or column 0 is kept empty. This makes the formulas short.

The grid used in this lesson

All examples use this grid, with 4 rows and 5 columns. The cell values are the numbers 1 to 20, written row by row.

       col 1  col 2  col 3  col 4  col 5
row 1    1      2      3      4      5
row 2    6      7      8      9     10
row 3   11     12     13     14     15
row 4   16     17     18     19     20

Take the rectangle from row 2 to row 3 and from column 2 to column 4. It contains the cells 7, 8, 9, 12, 13 and 14. Their sum is 63. The aim is to get this answer, and every other rectangle answer, without adding up the cells each time.

Part 1: the two-dimensional prefix table

The idea in plain words

A prefix sum is the total of everything from the start up to a point. In one dimension, a prefix sum table stores, for each position i, the sum of the values from position 1 to position i. Then the sum from l to r is the prefix at r minus the prefix at l - 1. You subtract the part you do not want.

In two dimensions, the table p stores, for each cell (i, j), the sum of every cell with row at most i and column at most j. That is the block in the top-left corner of the grid, ending at (i, j). The value p[i][j] is one number that summarises a whole corner. Once every corner is known, a rectangle can be found from four corners.

Building the table

Each entry of the table can be built from entries that are already known. The rule is:

p[i][j] = grid[i][j] + p[i - 1][j] + p[i][j - 1] - p[i - 1][j - 1]

The formula says: take the corner above, take the corner to the left, add the current cell, and remove the overlap. The overlap is the block that was counted twice. It is the corner at (i - 1, j - 1), which belongs to both the corner above and the corner to the left.

Here is the table for the grid above, with a row of zeros and a column of zeros added at the top and on the left, so that row 0 and column 0 exist and are empty. Check the cell p[2][3]. It is the sum of rows 1 and 2 and columns 1 to 3, which is 1 + 2 + 3 + 6 + 7 + 8 = 27. That value appears in row 2, column 3 of the table.

   0   0   0   0   0   0
   0   1   3   6  10  15
   0   7  16  27  40  55
   0  18  39  63  90 120
   0  34  72 114 160 210

The bottom-right cell is 210, which is the sum of the whole grid.

Answering a rectangle

Now take the rectangle from row r1 to r2 and column c1 to c2. The corner p[r2][c2] covers everything from the top-left corner down to (r2, c2). That includes some cells we do not want, namely the rows above r1 and the columns to the left of c1. The corners p[r1 - 1][c2] and p[r2][c1 - 1] cover those unwanted parts, but they overlap in the corner p[r1 - 1][c1 - 1]. So we subtract both unwanted corners and add back the overlap once:

sum = p[r2][c2] - p[r1 - 1][c2] - p[r2][c1 - 1] + p[r1 - 1][c1 - 1]

Apply it to the rectangle from row 2 to row 3 and column 2 to column 4:

p[3][4] is 90.
p[1][4] is 10.
p[3][1] is 18.
p[1][1] is 1.

The answer is 90 - 10 - 18 + 1 = 63, which matches the sum of the six cells.

Read the table as p[row][column], row first and column second. The entry p[3][1] is in row 3 and column 1, which is 18. It is the sum of the cells 1, 6 and 11 in column 1, down to row 3. Reading it as row 1, column 3 gives 6 instead, and the answer would be wrong.

The formula costs four list reads and some arithmetic, no matter how big the rectangle is. A rectangle of 1000 cells and a rectangle of one million cells take the same time.

The code for the prefix table

The following program builds the table, answers two rectangles from it and checks the first rectangle by adding the cells one by one. The Fenwick part is added in the second half of the lesson, and it is in the same file.

Python · editable

import random

def build_prefix(g, n, m):
    p = [[0] * (m + 1) for _ in range(n + 1)]
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            p[i][j] = g[i][j] + p[i - 1][j] + p[i][j - 1] - p[i - 1][j - 1]
    return p

def rect_from_prefix(p, r1, c1, r2, c2):
    return p[r2][c2] - p[r1 - 1][c2] - p[r2][c1 - 1] + p[r1 - 1][c1 - 1]

class Fenwick2D:
    def __init__(self, n, m):
        self.n = n
        self.m = m
        self.t = [[0] * (m + 1) for _ in range(n + 1)]

def add(self, r, c, delta):
        i = r
        while i <= self.n:
            j = c
            while j <= self.m:
                self.t[i][j] += delta
                j += j & -j
            i += i & -i

def prefix(self, r, c):
        s = 0
        i = r
        while i > 0:
            j = c
            while j > 0:
                s += self.t[i][j]
                j -= j & -j
            i -= i & -i
        return s

def rect(self, r1, c1, r2, c2):
        return self.prefix(r2, c2) - self.prefix(r1 - 1, c2) - self.prefix(r2, c1 - 1) + self.prefix(r1 - 1, c1 - 1)

def brute_rect(g, r1, c1, r2, c2):
    return sum(g[i][j] for i in range(r1, r2 + 1) for j in range(c1, c2 + 1))

n, m = 4, 5
grid = [[0] * (m + 1) for _ in range(n + 1)]
value = 1
for i in range(1, n + 1):
    for j in range(1, m + 1):
        grid[i][j] = value
        value += 1

p = build_prefix(grid, n, m)
print("prefix table:")
for row in p:
    print("".join(f"{v:4d}" for v in row))
print("rect (2,2)-(3,4) from prefix table:", rect_from_prefix(p, 2, 2, 3, 4))
print("rect (1,1)-(4,5) from prefix table:", rect_from_prefix(p, 1, 1, 4, 5))

fw = Fenwick2D(n, m)
for i in range(1, n + 1):
    for j in range(1, m + 1):
        fw.add(i, j, grid[i][j])
print("fenwick rect (2,2)-(3,4):", fw.rect(2, 2, 3, 4))
print("fenwick prefix(4,5):", fw.prefix(4, 5))

fw.add(3, 3, 100)
grid[3][3] += 100
print("after adding 100 to cell (3,3):")
print("fenwick rect (2,2)-(3,4):", fw.rect(2, 2, 3, 4))
print("fenwick rect (3,3)-(3,3):", fw.rect(3, 3, 3, 3))
print("stale prefix answer:", rect_from_prefix(p, 2, 2, 3, 4))

big_n, big_m = 60, 70
g = [[0] * (big_m + 1) for _ in range(big_n + 1)]
big = Fenwick2D(big_n, big_m)
rnd = random.Random(11)
ok = True
checks = 0
for op in range(3000):
    if rnd.random() < 0.5:
        r = 1 + rnd.randrange(big_n)
        c = 1 + rnd.randrange(big_m)
        d = rnd.randrange(2001) - 1000
        g[r][c] += d
        big.add(r, c, d)
    else:
        r1 = 1 + rnd.randrange(big_n)
        r2 = 1 + rnd.randrange(big_n)
        c1 = 1 + rnd.randrange(big_m)
        c2 = 1 + rnd.randrange(big_m)
        tr1, tr2 = min(r1, r2), max(r1, r2)
        tc1, tc2 = min(c1, c2), max(c1, c2)
        checks += 1
        if big.rect(tr1, tc1, tr2, tc2) != brute_rect(g, tr1, tc1, tr2, tc2):
            ok = False
print("random test:", checks, "rectangle checks, all match brute force:", ok)

The output is:

prefix table:
   0   0   0   0   0   0
   0   1   3   6  10  15
   0   7  16  27  40  55
   0  18  39  63  90 120
   0  34  72 114 160 210
rect (2,2)-(3,4) from prefix table: 63
rect (1,1)-(4,5) from prefix table: 210
fenwick rect (2,2)-(3,4): 63
fenwick prefix(4,5): 210
after adding 100 to cell (3,3):
fenwick rect (2,2)-(3,4): 163
fenwick rect (3,3)-(3,3): 113
stale prefix answer: 63
random test: 1479 rectangle checks, all match brute force: True

The prefix table lines match the table drawn earlier. The first rectangle gives 63, which is the answer worked out by hand. The whole grid gives 210, the bottom-right cell.

The stale line shows the weakness of the table. After the change to cell (3, 3), the old table still answers 63, but the correct answer is 163.

The random test uses a 60 by 70 grid. It runs 3000 operations, where each one is either a change to a random cell or a question about a random rectangle. It asked 1479 rectangle questions, and every answer matched the slow method that adds the cells one by one. The random generator is seeded with 11, so the same operations run on every run.

Python integers never overflow, so the totals in this program are exact no matter how large they grow. In a language with fixed-size integers, the same sums could wrap around, and the table would need a wider type. The concern is about the language and not about the formulas.

What the prefix table costs

Building the table takes time proportional to the number of cells, which is n * m steps. Each rectangle question takes four reads. So the table is very good when the grid is fixed and many questions are asked.

The table is weak when the grid changes. Changing one cell changes the corner sums of every cell below and to the right of it. In a grid of 1000 by 1000 cells, that can be up to a million table entries. Doing this after each change is too slow when changes and questions are mixed. The next part solves that.

Part 2: the two-dimensional Fenwick tree

A short recap of the one-dimensional Fenwick tree

A one-dimensional Fenwick tree, also called a binary indexed tree, stores a block of values in each slot. The lesson on counting inversions with a Fenwick tree, earlier in this module, explains it in full. Here is the rule in short form. Slot i stores the sum of the values from position i - lowbit(i) + 1 to position i, where lowbit(i) is the value of the lowest 1 bit of i. In Python, lowbit(i) is written i & -i.

For example, slot 6 is 110 in binary, and its lowest 1 bit is the 2 in position 010. So slot 6 covers the two positions 5 and 6. Slot 4 is 100 in binary, so its lowest 1 bit is 4, and it covers positions 1 to 4. Slot 3 is 011, its lowest 1 bit is 1, and it covers position 3 only.

Two rules follow from this. To add a value at position i, move to the next slot that covers i, which is i + lowbit(i), and repeat until the end. To get a prefix sum up to position i, add the slot i, then move to i - lowbit(i), and repeat until the start. Each walk takes about as many steps as there are bits in the number.

Extending the idea to two dimensions

In two dimensions, each slot is a pair (i, j). The slot t[i][j] covers the rows from i - lowbit(i) + 1 to i, and the columns from j - lowbit(j) + 1 to j. It is a rectangle, not a line, and its size is set by the lowest 1 bits of the row and column numbers.

To add a value at (r, c), the update walks up the rows and up the columns at the same time, using the same jumps as before. The update touches every slot whose rectangle contains (r, c).

To get a prefix sum up to (r, c), the query walks down the rows and down the columns at the same time. It adds the slot at each step, and that gives the sum of the whole corner from (1, 1) to (r, c).

The rectangle question then uses the same four-corner formula as the prefix table, with the prefix function in place of the table:

sum = prefix(r2, c2) - prefix(r1 - 1, c2) - prefix(r2, c1 - 1) + prefix(r1 - 1, c1 - 1)

A worked trace on the grid

Take the prefix sum up to row 3 and column 4, which should be 90, the value in the table at p[3][4].

The row walk starts at 3. Its lowest 1 bit is 1, so the next row is 3 - 1 = 2. The lowest 1 bit of 2 is 2, so the next row is 0 and the walk stops. The rows used are 3 and 2.

The column walk starts at 4. Its lowest 1 bit is 4, so the next column is 0 and the walk stops. The only column used is 4.

So the prefix is t[3][4] + t[2][4]. The slot t[3][4] covers row 3 and columns 1 to 4, which holds 11 + 12 + 13 + 14 = 50. The slot t[2][4] covers rows 1 and 2 and columns 1 to 4, which holds 1 + 2 + 3 + 4 + 6 + 7 + 8 + 9 = 40. Their sum is 90, which matches the table.

Now take an update. Add 100 to the cell (3, 3). The row walk starts at 3 and moves to 3 + 1 = 4, because the lowest 1 bit of 3 is 1. From 4 it moves to 8, which is past the 4 rows, so it stops. The rows touched are 3 and 4. The column walk works the same way, with columns 3 and 4. The four slots updated are t[3][3], t[3][4], t[4][3] and t[4][4]. Each of them covers the cell (3, 3), so each one gets the extra 100.

Those four slots are enough for later questions. The rectangle from (3, 3) to (3, 3) is prefix(3, 3) - prefix(2, 3) - prefix(3, 2) + prefix(2, 2), which the program computes as 113, the old value 13 plus 100.

Cost

For a grid of 1000 by 1000 cells, the number of bits in any row or column number is at most 10. Each walk takes at most about 10 steps, so one update or one prefix query touches at most about 100 slots. A rectangle question uses four prefix queries, so it touches at most about 400 slots. That is far less than the million cells the prefix table would need to rebuild after a change, and far less than the cells added one by one.

The price is that each answer is a little slower than the four reads of the prefix table, and the code is harder to read. Use the Fenwick tree when the grid changes and questions are mixed, and use the prefix table when the grid is fixed.

A second example: seats in a room

A seating chart is a grid where each cell is 0 for an empty seat and 1 for a taken seat. A person sits down, which adds 1 to a cell. A person stands up, which subtracts 1. A question asks how many seats are taken in a block. This is a rectangle question that changes all the time, so it suits the Fenwick tree.

The program below uses the Fenwick2D class from the program above. It runs six events and asks for one block after each event. It then adds up the plain grid as a check. Press Run to see the output.

Python · editable

class Fenwick2D:
    def __init__(self, n, m):
        self.n = n
        self.m = m
        self.t = [[0] * (m + 1) for _ in range(n + 1)]

def rect(self, r1, c1, r2, c2):
        return self.prefix(r2, c2) - self.prefix(r1 - 1, c2) - self.prefix(r2, c1 - 1) + self.prefix(r1 - 1, c1 - 1)

rows, cols = 5, 6
seats = [[0] * (cols + 1) for _ in range(rows + 1)]
room = Fenwick2D(rows, cols)
events = ["sit 2 3", "sit 2 4", "sit 3 3", "sit 5 6", "stand 2 3", "sit 1 1"]
for event in events:
    parts = event.split()
    r = int(parts[1])
    c = int(parts[2])
    change = 1 if parts[0] == "sit" else -1
    seats[r][c] += change
    room.add(r, c, change)
    print(event, "-> seats in rows 2 to 3, cols 3 to 4:", room.rect(2, 3, 3, 4))
print("whole room:", room.rect(1, 1, rows, cols))

total = 0
for row in seats:
    for v in row:
        total += v
print("check by summing the plain grid:", total)

The output is:

sit 2 3 -> seats in rows 2 to 3, cols 3 to 4: 1
sit 2 4 -> seats in rows 2 to 3, cols 3 to 4: 2
sit 3 3 -> seats in rows 2 to 3, cols 3 to 4: 3
sit 5 6 -> seats in rows 2 to 3, cols 3 to 4: 3
stand 2 3 -> seats in rows 2 to 3, cols 3 to 4: 2
sit 1 1 -> seats in rows 2 to 3, cols 3 to 4: 2
whole room: 4
check by summing the plain grid: 4

Read the lines one at a time. After the first event, the person at (2, 3) is in the block, so the count is 1. After the second event, the person at (2, 4) is in the block, so the count is 2. The event sit 5 6 is outside the block, so the block count stays at 3. When the person at (2, 3) stands up, the count drops to 2. The last event is outside the block. The whole room has four people, and the plain grid gives the same total.

The plain grid is kept only for this check. In a real program that needs both the seats and the counts, the grid is still useful, because it tells you the current state of one seat in one step.

Common mistakes

Mixing 0-based and 1-based indices. The formulas use 1-based indices, and Python lists use 0-based indices. The code makes the lists one larger so that index 0 is unused. If you size the list with the grid size, the last row or column is missing. Here is a short program that makes this mistake:

n = 4
p = [[0] * (n + 1) for _ in range(n)]
for row in range(1, n + 1):
    p[row][1] = row
print(p[4][1])

Python reports this error (the marker lines depend on the Python version):

Traceback (most recent call last):
  File "zerobased.py", line 4, in <module>
    p[row][1] = row
    ~^^^^^
IndexError: list index out of range

The list has rows 0 to 3, so row 4 does not exist. The caret points at the part of the line that failed, which is p[row]. The fix is to make the row count n + 1, as the program above does.

An update loop that starts at 0. The Fenwick loops depend on i & -i. For i = 0, that value is 0, so i never changes and the loop never ends. If the first position is 0, the loop runs forever. A check with a small loop shows the problem:

n = 4
t = [0] * (n + 1)
i = 0
while i <= n:
    t[i] += 1
    print("updated", i, flush=True)
    i += i & -i

The output repeats the same line without end, and the program never reaches its last line. The first three lines of its output are:

updated 0
updated 0
updated 0

Run this check only with a timeout, for example by pressing Ctrl+C after a second. Start all Fenwick positions at 1. If your data uses 0 for the first cell, add 1 to every row and column number before you call the Fenwick tree.

Forgetting to subtract the unwanted corners. The rectangle formula has four terms with the signs plus, minus, minus and plus. If you return only p[r2][c2], you get the whole corner, not the rectangle. For the rectangle from row 2 to row 3 and column 2 to column 4, that returns 90 instead of 63. The minus terms remove the rows above row 2 and the columns left of column 2. The plus term adds back the corner that was removed twice.

Using a stale prefix table after a change. A prefix table is a snapshot. If you change grid[3][3] by 100 and do not rebuild the table, the next rectangle answer is the old one. For the rectangle from row 2 to row 3 and column 2 to column 4, the stale table still answers 63, but the correct answer is 163. Either rebuild the table after every change, which is slow for big grids, or use the Fenwick tree.

Changing the cell but not the Fenwick tree, or the reverse. Keep the plain grid and the Fenwick tree in step. The program above adds the same change to both, and the random test compares the Fenwick answer with the plain grid. If you forget one of the two updates, the random test fails, and that is the easiest way to find the mistake.

Practice

Exercise 1. Use this 3 by 3 grid, with values 1 to 9 written row by row. Build the prefix table with a row of zeros and a column of zeros. Then answer the rectangle from row 2 to row 3 and column 2 to column 3, using only the four corners.

1  2  3
4  5  6
7  8  9

Show the answer

The table entries needed are p[3][3] = 45, p[1][3] = 6, p[3][1] = 12 and p[1][1] = 1. The sum is 45 - 6 - 12 + 1 = 28. Check it directly: the rectangle holds 5, 6, 8 and 9, and 5 + 6 + 8 + 9 = 28.

Exercise 2. In a Fenwick tree with 6 rows, which rows does the slot t[6] cover? Which rows does t[4] cover?

Show the answer

Slot 6 is 110 in binary. Its lowest 1 bit is 2, so it covers rows 5 and 6. Slot 4 is 100 in binary. Its lowest 1 bit is 4, so it covers rows 1 to 4. The rule is that slot i covers the rows from i - lowbit(i) + 1 to i.

Exercise 3. A Fenwick tree is built for a grid of 1000 rows and 1000 columns. What is the largest number of slots that one update can change?

Show the answer

The row walk and the column walk each take at most 10 steps, because the lowest 1 bit doubles at each step and the numbers are below 1024. The update changes the slots in every pair of walk positions, so the largest count is 10 times 10, which is 100 slots. The rule that each walk doubles its jump is what keeps the count small.

Exercise 4. A Fenwick tree with 4 rows and 4 columns gets add(2, 3, 5) and then add(4, 4, 2). What does rect(1, 1, 4, 4) return? What does rect(3, 3, 4, 4) return? Explain both answers.

Show the answer

rect(1, 1, 4, 4) returns 7, the sum of all the values added, because both changes are inside the whole grid. rect(3, 3, 4, 4) returns 2. The change at (2, 3) is in row 2, which is outside rows 3 and 4, so only the change at (4, 4) is counted.

Exercise 5. A grid of 500 by 500 cells gets 10000 changes and 10000 rectangle questions. Compare the total work for two plans. Plan A rebuilds a prefix table after every change. Plan B uses the Fenwick tree. Use only rough counts of table entries and slots.

Show the answer

For plan A, each rebuild touches all 250000 cells, so 10000 rebuilds touch 2.5 billion entries, plus 4 reads per question. For plan B, each walk takes at most about 9 steps for numbers below 512, so a change touches at most about 81 slots. A question uses four prefix queries, so it reads at most about 324 slots. Plan B does roughly 10000 times 81 for the changes, and 10000 times 324 for the questions, which is about 4 million slot reads in total. Plan B is far less work, which is why it is used for grids that change.

Summary

A two-dimensional prefix table answers any rectangle with four reads, but it must be rebuilt after each change.
A rectangle sum is the four-corner formula: add the bottom-right corner, subtract the two side corners, and add back the overlap.
A two-dimensional Fenwick tree stores a rectangle in each slot, and an update or a prefix query walks through about 10 slots per dimension for grids up to 1000 by 1000.
Use 1-based indices in the Fenwick tree, because index 0 makes the loop run forever.
Keep the plain grid and the Fenwick tree in step when a cell changes.

2D prefix sums and 2D Fenwick trees in Python

Why this matters

The grid used in this lesson

Part 1: the two-dimensional prefix table

The idea in plain words

Building the table

Answering a rectangle

The code for the prefix table

What the prefix table costs

Part 2: the two-dimensional Fenwick tree

A short recap of the one-dimensional Fenwick tree

Extending the idea to two dimensions

A worked trace on the grid

Cost

A second example: seats in a room

Common mistakes

Practice

Summary

Stuck? Ask the tutor