Two-Dimensional Prefix Sums
24 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)Sales, marks and seats often sit in a table. A report asks for the total of a block of rows and columns, many times. A table of corner totals turns each block query into four lookups.
Adding a block cell by cell
A table holds sales per store (rows) and month (columns). A report asks for the total of a block of stores and months. The direct way adds every cell in the block.
grid = [ [3, 1, 2, 4], [5, 2, 0, 1], [1, 6, 3, 2],]
def block_slow(r1, c1, r2, c2): return sum(sum(row[c1:c2 + 1]) for row in grid[r1:r2 + 1])
print(block_slow(1, 1, 2, 3))Example 1 · Slow block sums
grid = [ [3, 1, 2, 4], [5, 2, 0, 1], [1, 6, 3, 2],]
for r1, c1, r2, c2 in [(0, 0, 0, 3), (1, 1, 2, 3), (2, 2, 2, 2)]: print(sum(sum(row[c1:c2 + 1]) for row in grid[r1:r2 + 1]))10 14 3
What this code does
- Each query slices the rows, slices every row, and adds all the cells.
- The queries are a whole top row, a 2 by 3 block and a single cell.
One query costs O(rows * cols) in the worst case. The prefix idea from the last lesson has a two-dimensional form.
Build the table of corner totals
2D prefix sum is a table where P[i][j] holds the sum of all cells in rows 0 to i - 1 and columns 0 to j - 1. Any block of the grid can then be summed with four lookups.
Let P[i][j] be the sum of every cell above row i and left of column j. Each cell adds the total above and the total on the left, then subtracts the overlap counted twice.
grid = [ [3, 1, 2, 4], [5, 2, 0, 1], [1, 6, 3, 2],]
rows, cols = len(grid), len(grid[0])P = [[0] * (cols + 1) for _ in range(rows + 1)]for i in range(rows): for j in range(cols): P[i + 1][j + 1] = grid[i][j] + P[i][j + 1] + P[i + 1][j] - P[i][j]
print(P[rows][cols])Example 2 · The prefix table of a small grid
grid = [[1, 2, 3], [4, 5, 6]]rows, cols = len(grid), len(grid[0])P = [[0] * (cols + 1) for _ in range(rows + 1)]for i in range(rows): for j in range(cols): P[i + 1][j + 1] = grid[i][j] + P[i][j + 1] + P[i + 1][j] - P[i][j]for row in P: print(row)[0, 0, 0, 0] [0, 1, 3, 6] [0, 5, 12, 21]
What this code does
- The first row and first column are zeros, so no query needs a special case.
- The last value, 21, is the sum of the whole grid.
Example 3 · The shared rows trap
shared = [[0] * 3] * 2shared[0][0] = 1print(shared)
separate = [[0] * 3 for _ in range(2)]separate[0][0] = 1print(separate)[[1, 0, 0], [1, 0, 0]] [[1, 0, 0], [0, 0, 0]]
What this code does
[[0] * 3] * 2repeats the same row object twice, so a change shows in both rows.- The comprehension builds a new row on every turn. Always create a table this way.
P has one extra row and one extra column of zeros. Building costs O(rows * cols) time and space, once.
Any block with four lookups
Take the big rectangle from the corner to the block. Cut off the strip above and the strip on the left. Add back the corner that was cut twice.
P = [ [0, 0, 0, 0, 0], [0, 3, 4, 6, 10], [0, 8, 11, 13, 18], [0, 9, 18, 23, 30],]
def block(r1, c1, r2, c2): return P[r2 + 1][c2 + 1] - P[r1][c2 + 1] - P[r2 + 1][c1] + P[r1][c1]
print(block(1, 1, 2, 3))Example 4 · Four lookups agree with the slow way
grid = [ [3, 1, 2, 4], [5, 2, 0, 1], [1, 6, 3, 2],]
rows, cols = len(grid), len(grid[0])P = [[0] * (cols + 1) for _ in range(rows + 1)]for i in range(rows): for j in range(cols): P[i + 1][j + 1] = grid[i][j] + P[i][j + 1] + P[i + 1][j] - P[i][j]
def block(r1, c1, r2, c2): return P[r2 + 1][c2 + 1] - P[r1][c2 + 1] - P[r2 + 1][c1] + P[r1][c1]
for r1, c1, r2, c2 in [(0, 0, 2, 3), (1, 1, 2, 3), (0, 2, 1, 2)]: slow = sum(sum(row[c1:c2 + 1]) for row in grid[r1:r2 + 1]) print(block(r1, c1, r2, c2), slow)30 30 14 14 2 2
What this code does
- Both columns match, which is a quick way to test your formula.
- The first query covers the whole grid and the third one is a column of two cells.
Four lookups and three arithmetic operations, so each query is O(1) whatever the size of the block.
Wrap it in a small class
Interviews often ask for a class that is built once and queried many times. The constructor does the O(rows * cols) work, and the method does the O(1) work.
class GridSum: def __init__(self, grid): rows, cols = len(grid), len(grid[0]) self.p = [[0] * (cols + 1) for _ in range(rows + 1)] for i in range(rows): for j in range(cols): self.p[i + 1][j + 1] = ( grid[i][j] + self.p[i][j + 1] + self.p[i + 1][j] - self.p[i][j] )
def block(self, r1, c1, r2, c2): p = self.p return p[r2 + 1][c2 + 1] - p[r1][c2 + 1] - p[r2 + 1][c1] + p[r1][c1]Example 5 · GridSum used on a sales table
class GridSum: def __init__(self, grid): rows, cols = len(grid), len(grid[0]) self.p = [[0] * (cols + 1) for _ in range(rows + 1)] for i in range(rows): for j in range(cols): self.p[i + 1][j + 1] = ( grid[i][j] + self.p[i][j + 1] + self.p[i + 1][j] - self.p[i][j] )
def block(self, r1, c1, r2, c2): p = self.p return p[r2 + 1][c2 + 1] - p[r1][c2 + 1] - p[r2 + 1][c1] + p[r1][c1]
sales = [[3, 1, 2, 4], [5, 2, 0, 1], [1, 6, 3, 2]]table = GridSum(sales)print(table.block(0, 0, 1, 1))print(table.block(1, 1, 2, 3))print(table.block(2, 0, 2, 3))11 14 12
What this code does
- The object is built once and then answers each block in O(1).
- The last query is the whole bottom row: 1 + 6 + 3 + 2.
Example 6 · The best 2 by 2 block
class GridSum: def __init__(self, grid): rows, cols = len(grid), len(grid[0]) self.p = [[0] * (cols + 1) for _ in range(rows + 1)] for i in range(rows): for j in range(cols): self.p[i + 1][j + 1] = ( grid[i][j] + self.p[i][j + 1] + self.p[i + 1][j] - self.p[i][j] )
def block(self, r1, c1, r2, c2): p = self.p return p[r2 + 1][c2 + 1] - p[r1][c2 + 1] - p[r2 + 1][c1] + p[r1][c1]
grid = [[3, 1, 2, 4], [5, 2, 0, 1], [1, 6, 3, 2]]table = GridSum(grid)best = max( (table.block(r, c, r + 1, c + 1), r, c) for r in range(len(grid) - 1) for c in range(len(grid[0]) - 1))print(best)(14, 1, 0)
What this code does
- The generator tries every top-left corner of a 2 by 2 block.
- Tuples compare from the left, so max picks the largest sum first. The result is (sum, row, column).
The grid is not stored. Only the prefix table is kept, and it is enough to answer every block query.
Next step
Check your understanding
12 quick questions