DSA Trapping Rain Water II

The problem

A height map is a grid where each cell holds a ground height. Trapping Rain Water II asks how much water the grid holds after rain. Water rests on a cell only when walls surround it on every path to the border.

Water escapes over the lowest wall on its best escape path. A single low border cell can drain a whole bowl. The cases below are the ones this article solves.

sealed = [[3, 3, 3], [3, 1, 3], [3, 3, 3]]
bowl = [[3, 3, 3, 3], [3, 1, 2, 3], [3, 2, 1, 3], [3, 3, 3, 3]]
leaky = [[3, 3, 3], [3, 1, 1], [3, 3, 3]]
classic = [[1, 4, 3, 1, 3, 2], [3, 2, 1, 3, 2, 4], [2, 3, 3, 2, 3, 1]]
for name, grid, expected in [("sealed", sealed, 2), ("bowl", bowl, 6), ("leaky", leaky, 0), ("classic", classic, 4)]:
    print(name, "grid rows:", len(grid), "water held:", expected)

The sealed and leaky bowls differ in one border cell. That one cell decides between two units of water and none.

Intuition: one wall height for the whole grid

The first idea gives every inner cell the same wall. The lowest border height looks like the level water can reach before spilling. A cell then holds that level minus its own height.

grid = [[2, 5, 5, 5], [5, 1, 5, 5], [5, 5, 5, 5], [5, 5, 5, 5]]
rows = len(grid)
cols = len(grid[0])
border = [grid[r][c] for r in range(rows) for c in range(cols) if r in (0, rows - 1) or c in (0, cols - 1)]
lowest_border = min(border)
cell = grid[1][1]
print("lowest border height:", lowest_border)
print("inner cell height:", cell)
print("single-wall guess for the cell:", lowest_border - cell)
neighbours = [grid[0][1], grid[2][1], grid[1][0], grid[1][2]]
print("the four walls around the cell:", neighbours)

The guess says one unit. Every wall around the cell is 5 high, so the cell truly holds four units. The low border cell sits in a corner the water never reaches. A wall is a property of a path, not of the whole grid.

Brute force: lower the levels until nothing changes

The brute force gives each inner cell a water level that starts at infinity. Border cells keep their own height as their level, because water leaves the grid there. A pass lowers each inner level to the lowest level around it, but never below the cell's own height.

Each pass reads the levels left by the previous pass. Certain levels spread one ring inward per pass. Passes repeat until a full pass changes nothing, and a counter records them.

def trap_slow(height_map):
    rows = len(height_map)
    cols = len(height_map[0])
    level = []
    for r in range(rows):
        row = []
        for c in range(cols):
            if r in (0, rows - 1) or c in (0, cols - 1):
                row.append(height_map[r][c])
            else:
                row.append(float("inf"))
        level.append(row)
    passes = 0
    while True:
        passes += 1
        changed = False
        settled = [current[:] for current in level]
        for r in range(1, rows - 1):
            for c in range(1, cols - 1):
                lowest = min(level[r - 1][c], level[r + 1][c], level[r][c - 1], level[r][c + 1])
                new_level = max(height_map[r][c], lowest)
                if new_level < level[r][c]:
                    settled[r][c] = new_level
                    changed = True
        level = settled
        if not changed:
            break
    water = 0
    for r in range(1, rows - 1):
        for c in range(1, cols - 1):
            water += level[r][c] - height_map[r][c]
    return water, passes

sealed = [[3, 3, 3], [3, 1, 3], [3, 3, 3]]
bowl = [[3, 3, 3, 3], [3, 1, 2, 3], [3, 2, 1, 3], [3, 3, 3, 3]]
leaky = [[3, 3, 3], [3, 1, 1], [3, 3, 3]]
classic = [[1, 4, 3, 1, 3, 2], [3, 2, 1, 3, 2, 4], [2, 3, 3, 2, 3, 1]]
walled = [[2, 5, 5, 5], [5, 1, 5, 5], [5, 5, 5, 5], [5, 5, 5, 5]]
for name, grid in [("sealed", sealed), ("bowl", bowl), ("leaky", leaky), ("classic", classic), ("walled", walled)]:
    water, passes = trap_slow(grid)
    print(name, "water:", water, "passes:", passes)

The answers are all correct, including 4 for the walled grid from the intuition. The cost is the repetition. A deep bowl settles one ring per pass, so large grids pay many full sweeps.

The optimal solution: flood inward from the lowest wall

A priority queue is a heap where the smallest entry leaves first, and flooding inward from the border always through the lowest current wall fixes each cell's level the first time the flood reaches it. That one rule is the whole method.

Every border cell starts in the heap and marked visited. The lowest wall pops out, and each unvisited neighbour traps the gap between that wall level and its height. The neighbour then joins the heap carrying the higher of the wall level and its own height.

import heapq


def trap(height_map):
    rows = len(height_map)
    cols = len(height_map[0])
    visited = [[False] * cols for _ in range(rows)]
    heap = []
    for r in range(rows):
        for c in range(cols):
            if r in (0, rows - 1) or c in (0, cols - 1):
                heapq.heappush(heap, (height_map[r][c], r, c))
                visited[r][c] = True
    water = 0
    while heap:
        level, r, c = heapq.heappop(heap)
        for dr, dc in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
            nr, nc = r + dr, c + dc
            if 0 <= nr < rows and 0 <= nc < cols and not visited[nr][nc]:
                visited[nr][nc] = True
                if height_map[nr][nc] < level:
                    water += level - height_map[nr][nc]
                wall = max(level, height_map[nr][nc])
                heapq.heappush(heap, (wall, nr, nc))
    return water

sealed = [[3, 3, 3], [3, 1, 3], [3, 3, 3]]
bowl = [[3, 3, 3, 3], [3, 1, 2, 3], [3, 2, 1, 3], [3, 3, 3, 3]]
leaky = [[3, 3, 3], [3, 1, 1], [3, 3, 3]]
classic = [[1, 4, 3, 1, 3, 2], [3, 2, 1, 3, 2, 4], [2, 3, 3, 2, 3, 1]]
walled = [[2, 5, 5, 5], [5, 1, 5, 5], [5, 5, 5, 5], [5, 5, 5, 5]]
for name, grid, expected in [("sealed", sealed, 2), ("bowl", bowl, 6), ("leaky", leaky, 0), ("classic", classic, 4), ("walled", walled, 4)]:
    result = trap(grid)
    print(name, "water:", result, "expected:", expected, "match:", result == expected)

All five grids match. The leaky bowl drains because its low border cell pops first and reaches the centre at level 1. No cell is ever revisited once marked.

Walkthrough: the flood on the sealed bowl

The sealed bowl is three by three with a centre of height 1. Eight border cells of height 3 enter the heap first. The trace prints each pop and the one discovery that matters.

import heapq

height_map = [[3, 3, 3], [3, 1, 3], [3, 3, 3]]
rows = 3
cols = 3
visited = [[False] * cols for _ in range(rows)]
heap = []
for r in range(rows):
    for c in range(cols):
        if r in (0, rows - 1) or c in (0, cols - 1):
            heapq.heappush(heap, (height_map[r][c], r, c))
            visited[r][c] = True
print("border cells pushed:", len(heap))
water = 0
while heap:
    level, r, c = heapq.heappop(heap)
    print("pop wall level", level, "at", (r, c))
    for dr, dc in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
        nr, nc = r + dr, c + dc
        if 0 <= nr < rows and 0 <= nc < cols and not visited[nr][nc]:
            visited[nr][nc] = True
            held = level - height_map[nr][nc]
            water += held
            print("  cell", (nr, nc), "height", height_map[nr][nc], "traps", held)
            heapq.heappush(heap, (max(level, height_map[nr][nc]), nr, nc))
print("total water:", water)

The second pop reaches the centre at level 3, and the centre traps 2 at once. Its level is final in that moment, because no lower wall can ever reach it later. The heap pops lowest first, so any later path would arrive carrying a wall at least as high.

Complexity

The heap flood pushes each cell exactly once and pops it exactly once, so its time is O(N log N) for N cells. It stores the heap and the visited grid, so its extra space is O(N).

The settling process sweeps every inner cell per pass, and a cone-shaped bowl needs one pass per ring plus a final confirming pass. The counts below build cone bowls of growing size and compare cell visits.

def cone(size):
    grid = []
    for r in range(size):
        row = []
        for c in range(size):
            depth = min(r, c, size - 1 - r, size - 1 - c)
            row.append(size // 2 + 1 - depth)
        grid.append(row)
    return grid


def slow_visits(height_map):
    rows = len(height_map)
    cols = len(height_map[0])
    level = []
    for r in range(rows):
        row = []
        for c in range(cols):
            if r in (0, rows - 1) or c in (0, cols - 1):
                row.append(height_map[r][c])
            else:
                row.append(float("inf"))
        level.append(row)
    visits = 0
    while True:
        changed = False
        settled = [current[:] for current in level]
        for r in range(1, rows - 1):
            for c in range(1, cols - 1):
                visits += 1
                lowest = min(level[r - 1][c], level[r + 1][c], level[r][c - 1], level[r][c + 1])
                new_level = max(height_map[r][c], lowest)
                if new_level < level[r][c]:
                    settled[r][c] = new_level
                    changed = True
        level = settled
        if not changed:
            return visits

for size in [3, 5, 9, 17]:
    grid = cone(size)
    cells = size * size
    print("grid:", size, "x", size, "- cells:", cells, " heap pushes:", cells, " settling visits:", slow_visits(grid))

At 17 by 17 the settling process visits 2,025 cells to do what 289 pushes do. The flood pays a small logarithmic price per cell and never sweeps twice.

Edge cases and pitfalls

A grid with fewer than three rows or columns has no inner cells, so it holds nothing. A flat grid holds nothing, because every cell already sits at wall level. A tall inner cell holds nothing either, because water runs off a peak.

import heapq


def trap(height_map):
    rows = len(height_map)
    cols = len(height_map[0])
    visited = [[False] * cols for _ in range(rows)]
    heap = []
    for r in range(rows):
        for c in range(cols):
            if r in (0, rows - 1) or c in (0, cols - 1):
                heapq.heappush(heap, (height_map[r][c], r, c))
                visited[r][c] = True
    water = 0
    while heap:
        level, r, c = heapq.heappop(heap)
        for dr, dc in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
            nr, nc = r + dr, c + dc
            if 0 <= nr < rows and 0 <= nc < cols and not visited[nr][nc]:
                visited[nr][nc] = True
                if height_map[nr][nc] < level:
                    water += level - height_map[nr][nc]
                wall = max(level, height_map[nr][nc])
                heapq.heappush(heap, (wall, nr, nc))
    return water

print("single row:", trap([[1, 2, 3]]))
print("two rows:", trap([[3, 3, 3], [3, 3, 3]]))
print("flat grid:", trap([[2, 2, 2], [2, 2, 2], [2, 2, 2]]))
print("tall centre:", trap([[2, 2, 2], [2, 9, 2], [2, 2, 2]]))
print("empty centre at zero:", trap([[1, 1, 1], [1, 0, 1], [1, 1, 1]]))

One pitfall sits in the visited marks. A cell must be marked when it is pushed, not when it pops. Marking at pop time lets the same cell enter the heap through several neighbours, and a later entry can carry a higher wall and count its water twice.