LeetCode 407: Trapping Rain Water II

Problem Restatement

We are given an m x n matrix called heightMap.

Each cell represents the height of one unit block in a 2D elevation map.

After rain, water may be trapped in lower cells if surrounding boundaries prevent it from flowing out.

Return the total amount of trapped water.

Water can move through the four main directions:

up, down, left, right

Diagonal movement does not matter.

Boundary cells cannot trap water by themselves because water can flow out of the map from the boundary.

The constraints include 1 <= m, n <= 200 and 0 <= heightMap[i][j] <= 2 * 10^4. The standard examples include [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]] -> 4 and [[3,3,3,3,3],[3,2,2,2,3],[3,2,1,2,3],[3,2,2,2,3],[3,3,3,3,3]] -> 10.

Input and Output

Example function shape:

def trapRainWater(heightMap: list[list[int]]) -> int:
    ...

Examples

Example 1:

heightMap = [
    [1, 4, 3, 1, 3, 2],
    [3, 2, 1, 3, 2, 4],
    [2, 3, 3, 2, 3, 1],
]

The answer is:

4

The low inner cells can hold water, but the amount depends on the lowest boundary through which water could escape.

Example 2:

heightMap = [
    [3, 3, 3, 3, 3],
    [3, 2, 2, 2, 3],
    [3, 2, 1, 2, 3],
    [3, 2, 2, 2, 3],
    [3, 3, 3, 3, 3],
]

The answer is:

10

The outer ring has height 3, so inner cells lower than 3 can hold water up to level 3.

First Thought: Extend 1D Trapping Rain Water

In the 1D version, the water above an index depends on the maximum height to the left and the maximum height to the right.

For 2D, this becomes harder.

A cell can leak through many possible paths to the boundary. Its water level depends on the lowest boundary path, not only the tallest wall in four directions.

So a simple row-by-row or column-by-column method fails.

We need to process the map from the outside inward.

Key Insight

Water escapes through the boundary.

So the boundary determines the first water level.

The lowest boundary cell is the weakest wall. Any nearby lower cell can hold water only up to that lowest boundary height.

This suggests a min heap.

We start with all boundary cells in a min heap. Then we repeatedly expand from the current lowest boundary.

When we visit an unvisited neighbor:

If the neighbor is lower than the current boundary height, it traps:

current_boundary_height - neighbor_height

water.

Then the neighbor becomes part of the boundary for future expansion. Its effective height is:

max(current_boundary_height, neighbor_height)

because if it was filled with water, the new water surface is at least the current boundary height.

Algorithm

Use a priority queue ordered by height.

1. Push every boundary cell into the heap.
2. Mark every boundary cell as visited.
3. Initialize water = 0.
4. While the heap is not empty:
   • Pop the cell with the smallest effective height.
   • Check its four neighbors.
   • For each unvisited neighbor:
     • Mark it visited.
     • If the neighbor is lower, add trapped water.
     • Push the neighbor back with effective height max(current_height, neighbor_height).
5. Return water.

Correctness

The algorithm always expands from the lowest known boundary.

At any moment, the heap contains the boundary between the visited outside region and the unvisited inside region. The smallest heap cell is the lowest escape route available from the outside.

When the algorithm visits a neighbor from this lowest boundary, there are two cases.

If the neighbor height is lower than the boundary height, water can be trapped above it. The trapped amount is exactly the difference between the boundary height and the neighbor height. The cell cannot hold more than this level because the current boundary is already an escape path.

If the neighbor height is greater than or equal to the boundary height, it traps no water and becomes a higher wall for later cells.

The effective height pushed into the heap is max(current_height, neighbor_height). This records the lowest water level or wall height that future inner cells must overcome to escape.

Because cells are processed in increasing boundary height order, the first time a cell is visited, its minimum possible escape boundary has already been determined. Later paths cannot produce a lower boundary because the heap would have processed that lower boundary first.

Therefore, each cell contributes exactly the amount of water it can trap, and the final total is correct.

Complexity

Here m is the number of rows and n is the number of columns.

Implementation

import heapq
from typing import List

class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        m = len(heightMap)
        n = len(heightMap[0])

        if m < 3 or n < 3:
            return 0

        heap = []
        visited = [[False] * n for _ in range(m)]

        for r in range(m):
            for c in (0, n - 1):
                heapq.heappush(heap, (heightMap[r][c], r, c))
                visited[r][c] = True

        for c in range(1, n - 1):
            for r in (0, m - 1):
                heapq.heappush(heap, (heightMap[r][c], r, c))
                visited[r][c] = True

        water = 0
        directions = [(1, 0), (-1, 0), (0, 1), (0, -1)]

        while heap:
            height, r, c = heapq.heappop(heap)

            for dr, dc in directions:
                nr = r + dr
                nc = c + dc

                if nr < 0 or nr >= m or nc < 0 or nc >= n:
                    continue

                if visited[nr][nc]:
                    continue

                visited[nr][nc] = True

                neighbor_height = heightMap[nr][nc]

                if neighbor_height < height:
                    water += height - neighbor_height

                heapq.heappush(
                    heap,
                    (max(height, neighbor_height), nr, nc),
                )

        return water

Code Explanation

We first get the matrix size:

m = len(heightMap)
n = len(heightMap[0])

If either dimension is less than 3, there is no enclosed interior:

if m < 3 or n < 3:
    return 0

Then we create a min heap and a visited matrix:

heap = []
visited = [[False] * n for _ in range(m)]

All boundary cells are pushed into the heap.

First, the left and right columns:

for r in range(m):
    for c in (0, n - 1):
        heapq.heappush(heap, (heightMap[r][c], r, c))
        visited[r][c] = True

Then, the top and bottom rows, skipping corners to avoid duplicates:

for c in range(1, n - 1):
    for r in (0, m - 1):
        heapq.heappush(heap, (heightMap[r][c], r, c))
        visited[r][c] = True

The heap stores:

(effective_height, row, col)

The main loop always expands from the lowest effective boundary:

height, r, c = heapq.heappop(heap)

For every unvisited neighbor, we compare its ground height with the current boundary height:

if neighbor_height < height:
    water += height - neighbor_height

Then we push it with the new effective height:

max(height, neighbor_height)

This value means:

If the neighbor was lower, it is now filled up to height.

If the neighbor was higher, it becomes the new wall height.

Testing

def test_trap_rain_water():
    s = Solution()

    assert s.trapRainWater([
        [1, 4, 3, 1, 3, 2],
        [3, 2, 1, 3, 2, 4],
        [2, 3, 3, 2, 3, 1],
    ]) == 4

    assert s.trapRainWater([
        [3, 3, 3, 3, 3],
        [3, 2, 2, 2, 3],
        [3, 2, 1, 2, 3],
        [3, 2, 2, 2, 3],
        [3, 3, 3, 3, 3],
    ]) == 10

    assert s.trapRainWater([[1, 2, 3]]) == 0

    assert s.trapRainWater([
        [1, 1, 1],
        [1, 0, 1],
        [1, 1, 1],
    ]) == 1

    assert s.trapRainWater([
        [5, 5, 5, 5],
        [5, 1, 1, 5],
        [5, 5, 5, 5],
    ]) == 8

    print("all tests passed")

Test Notes