Trapping Rain Water II

Trapping Rain Water II takes the classic 1-D rainwater problem into two dimensions. The 1-D version scans with two pointers, but a 2-D grid has no single direction to sweep — water can escape in any direction. The fix is a min-heap that always processes the lowest wall first.

The slow way first

You might try to compute, for every inner cell, the lowest wall on the path to the border. But in 2-D the "lowest enclosing wall" depends on every possible escape route, not just four directions. Brute-forcing that — flooding from each cell and tracking the minimum barrier — is hopelessly slow, easily O((mn)²) or worse.

The question to ask: which cell can I resolve with certainty right now? The cell sitting behind the globally lowest wall on the current boundary — because no lower wall can ever surround it.

The idea: always flood from the lowest wall

Treat the border as the starting water boundary and push every border cell into a min-heap keyed by height. Then repeatedly pop the lowest wall. For each inner neighbor: the water level there is capped by this wall, so it traps wall − cellHeight (if positive). Push the neighbor back with its effective height max(wall, cellHeight), because from now on that is the lowest barrier protecting cells further inside.

Because the heap always serves the lowest boundary wall, the first time we reach any inner cell we already know the true lowest wall surrounding it — so we visit each cell exactly once.

Walk through it

Step through the animation. The border cells turn visited and enter the heap. We pop the lowest wall (a 2), look at its neighbor 1, and trap 2 − 1 = 1 unit. That neighbor re-enters the heap at effective height 2. Popping it next reaches a 2-cell where no water pools. When the heap empties, the running total is the answer.

Pseudocode

push every border cell into a min-heap keyed by height; mark them seen
water = 0
while the heap is not empty:
    (h, r, c) = pop the lowest wall
    for each in-grid neighbor (nr, nc) not yet seen:
        mark (nr, nc) seen
        water += max(0, h - height[nr][nc])      # trapped if wall is taller
        push (max(h, height[nr][nc]), nr, nc)     # effective new wall height
return water

The Python solution

def trap_rain_water(heights):
    R, C = len(heights), len(heights[0])
    heap, seen = [], set()
    for r in range(R):
        for c in range(C):
            if r in (0, R-1) or c in (0, C-1):
                heapq.heappush(heap, (heights[r][c], r, c)); seen.add((r, c))
    water = 0
    while heap:
        h, r, c = heapq.heappop(heap)
        for nr, nc in neighbors(r, c):
            if (nr, nc) not in seen:
                seen.add((nr, nc))
                water += max(0, h - heights[nr][nc])
                heapq.heappush(heap, (max(h, heights[nr][nc]), nr, nc))
    return water

• We seed the heap with the whole border, since border cells can never hold water.
• heapq.heappop always returns the lowest wall still on the boundary — that ordering is the whole trick.
• For each unseen neighbor, max(0, h - heights[nr][nc]) is the water trapped: only positive when the wall h is taller than the cell.
• We push the neighbor back with max(h, heights[nr][nc]) — its effective height as the new barrier for cells deeper inside.
• seen guarantees every cell is processed exactly once.

Complexity

Every cell enters and leaves the heap once, and each heap operation costs O(log(mn)), giving O(mn log(mn)). The heap plus the seen set use O(mn) space.

When this pattern shows up

Practice