LeetCode 407 - Trapping Rain Water II

This problem is the two dimensional version of the classic "Trapping Rain Water" problem. Instead of a one dimensional array of heights, we are given an m x n grid where each cell represents the elevation of a block in a terrain.

LeetCode Problem 407

Difficulty: 🔴 Hard
Topics: Array, Breadth-First Search, Heap (Priority Queue), Matrix

Solution

LeetCode 407 - Trapping Rain Water II

Problem Understanding

This problem is the two dimensional version of the classic "Trapping Rain Water" problem. Instead of a one dimensional array of heights, we are given an m x n grid where each cell represents the elevation of a block in a terrain.

The goal is to determine how much water can remain trapped after rainfall.

Each cell has a height, and water can only remain trapped if it is surrounded by taller boundaries. If there is any path from a cell to the outer border through lower or equal heights, water can escape and cannot be trapped there.

The input is:

  • heightMap[i][j], the height of the cell at row i and column j
  • m rows and n columns

The output is a single integer representing the total amount of trapped water.

The constraints are important:

  • 1 <= m, n <= 200
  • Heights can be as large as 2 * 10^4

A grid of size 200 x 200 contains up to 40,000 cells. Any solution worse than roughly O(m * n * log(m * n)) risks timing out. This immediately rules out many naive repeated simulation approaches.

There are several important observations and edge cases:

  • Border cells can never trap water because water can always flow out of the map from the boundary.
  • Very small grids, such as 1 x n, m x 1, or even 2 x 2, cannot trap water because there is no enclosed interior region.
  • Water level at a cell depends not only on its immediate neighbors but on the minimum boundary surrounding the entire connected region.
  • A locally tall wall does not guarantee trapped water if another path to the boundary exists through a lower wall.

The main challenge is efficiently determining the effective boundary height for every interior cell.

Approaches

Brute Force Approach

A brute force solution would examine every interior cell independently.

For each cell, we could try to determine the tallest water level that can remain there without leaking to the border. One way to think about this is:

  • Start from the current cell.
  • Explore outward in all directions.
  • Determine the minimum boundary height surrounding the cell.
  • The trapped water equals:

$$\text{boundary height} - \text{cell height}$$

if positive.

However, this is difficult because the surrounding boundary is global, not local. To correctly determine whether water escapes, we would effectively need a search from each cell to the border.

If we perform BFS or DFS from every cell, the complexity becomes extremely large:

  • Up to 40,000 cells
  • Each search may visit most of the grid

This produces approximately O((m * n)^2) complexity, which is too slow.

Although the brute force idea is conceptually correct, it repeatedly recomputes the same information.

Optimal Approach, Min Heap + BFS

The key insight is that trapped water is determined by the lowest boundary surrounding a region.

This is extremely similar to how water fills a basin in the real world:

  • Water starts leaking from the lowest boundary.
  • The smallest surrounding wall determines the maximum water level.

Instead of processing cells independently, we process the terrain from the outside inward.

The strategy is:

  • Insert all border cells into a min heap.

  • Always expand from the currently lowest boundary.

  • When visiting a neighboring cell:

  • If the neighbor is lower than the current boundary height, water is trapped.

  • Otherwise, no water is trapped.

  • The neighbor then becomes part of the boundary with height:

$$\max(\text{current boundary height}, \text{neighbor height})$$

This works because once a cell is reached from the minimum possible boundary, we already know the smallest wall that can contain water around it.

This is essentially Dijkstra's algorithm applied to water levels.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O((m × n)^2) O(m × n) Recomputes reachable boundaries for each cell independently
Optimal O(m × n × log(m × n)) O(m × n) Uses min heap BFS from the boundary inward

Algorithm Walkthrough

  1. Initialize dimensions and handle small grids.

If the grid has fewer than 3 rows or columns, no water can be trapped because there is no enclosed interior region. 2. Create a visited matrix.

We must ensure each cell is processed exactly once. Without this, the same cell could be inserted into the heap multiple times. 3. Create a min heap.

The heap stores cells in the format:

(height, row, column)

The heap always gives us the currently lowest boundary cell. 4. Push all border cells into the heap.

Every outer boundary cell is added because borders define where water can escape.

These cells are also marked as visited immediately. 5. Start BFS using the heap.

Repeatedly pop the smallest height cell from the heap.

This cell represents the current minimum boundary height for expansion. 6. Explore the four neighboring cells.

For each neighbor:

  • Skip if out of bounds
  • Skip if already visited
  • Otherwise mark it visited
  1. Compute trapped water.

If the neighbor's height is smaller than the current boundary height:

$$\text{water trapped} = \text{boundary height} - \text{neighbor height}$$

Add this amount to the answer. 8. Update the effective boundary height.

The neighbor becomes part of the boundary for future expansion.

Its effective height is:

$$\max(\text{current boundary height}, \text{neighbor height})$$

This is critical because once water fills a lower cell, the water surface acts like a taller boundary. 9. Continue until the heap becomes empty.

At that point, every reachable cell has been processed and all trapped water has been counted.

Why it works

The algorithm maintains a very important invariant:

Every cell popped from the heap is processed using the minimum possible boundary height that can reach it.

Because the heap always expands the lowest boundary first, we never miss a smaller escape route later. This guarantees that the first time we process a cell, we already know the correct water level for that cell.

This is exactly the same correctness principle used in Dijkstra's shortest path algorithm.

Python Solution

from typing import List
import heapq

class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        if not heightMap or not heightMap[0]:
            return 0

        rows = len(heightMap)
        cols = len(heightMap[0])

        if rows < 3 or cols < 3:
            return 0

        visited = [[False] * cols for _ in range(rows)]
        min_heap = []

        # Add all border cells to the heap
        for row in range(rows):
            for col in [0, cols - 1]:
                heapq.heappush(min_heap, (heightMap[row][col], row, col))
                visited[row][col] = True

        for col in range(cols):
            for row in [0, rows - 1]:
                if not visited[row][col]:
                    heapq.heappush(min_heap, (heightMap[row][col], row, col))
                    visited[row][col] = True

        total_water = 0

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

        while min_heap:
            current_height, row, col = heapq.heappop(min_heap)

            for dr, dc in directions:
                new_row = row + dr
                new_col = col + dc

                if (
                    new_row < 0 or
                    new_row >= rows or
                    new_col < 0 or
                    new_col >= cols or
                    visited[new_row][new_col]
                ):
                    continue

                visited[new_row][new_col] = True

                neighbor_height = heightMap[new_row][new_col]

                if neighbor_height < current_height:
                    total_water += current_height - neighbor_height

                effective_height = max(current_height, neighbor_height)

                heapq.heappush(
                    min_heap,
                    (effective_height, new_row, new_col)
                )

        return total_water

The implementation follows the algorithm directly.

The first section validates the input and handles small grids that cannot trap water.

Next, the solution initializes a visited matrix and a min heap. Every border cell is inserted into the heap because borders form the initial escape boundary for water.

The BFS loop repeatedly extracts the lowest boundary cell. This guarantees that we always expand from the smallest currently known wall.

For each unvisited neighbor:

  • If the neighbor is lower, trapped water is added.
  • The neighbor is inserted back into the heap using the effective boundary height.

The effective height is crucial. Once water fills a cell, the water surface becomes the new boundary level for future exploration.

The algorithm processes every cell exactly once, making it efficient enough for the constraints.

Go Solution

package main

import (
	"container/heap"
)

type Cell struct {
	height int
	row    int
	col    int
}

type MinHeap []Cell

func (h MinHeap) Len() int {
	return len(h)
}

func (h MinHeap) Less(i, j int) bool {
	return h[i].height < h[j].height
}

func (h MinHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *MinHeap) Push(x interface{}) {
	*h = append(*h, x.(Cell))
}

func (h *MinHeap) Pop() interface{} {
	old := *h
	n := len(old)

	item := old[n-1]
	*h = old[:n-1]

	return item
}

func trapRainWater(heightMap [][]int) int {
	if len(heightMap) == 0 || len(heightMap[0]) == 0 {
		return 0
	}

	rows := len(heightMap)
	cols := len(heightMap[0])

	if rows < 3 || cols < 3 {
		return 0
	}

	visited := make([][]bool, rows)
	for i := range visited {
		visited[i] = make([]bool, cols)
	}

	minHeap := &MinHeap{}
	heap.Init(minHeap)

	// Add border cells
	for row := 0; row < rows; row++ {
		heap.Push(minHeap, Cell{heightMap[row][0], row, 0})
		heap.Push(minHeap, Cell{heightMap[row][cols-1], row, cols - 1})

		visited[row][0] = true
		visited[row][cols-1] = true
	}

	for col := 0; col < cols; col++ {
		if !visited[0][col] {
			heap.Push(minHeap, Cell{heightMap[0][col], 0, col})
			visited[0][col] = true
		}

		if !visited[rows-1][col] {
			heap.Push(minHeap, Cell{heightMap[rows-1][col], rows - 1, col})
			visited[rows-1][col] = true
		}
	}

	totalWater := 0

	directions := [][]int{
		{1, 0},
		{-1, 0},
		{0, 1},
		{0, -1},
	}

	for minHeap.Len() > 0 {
		current := heap.Pop(minHeap).(Cell)

		for _, dir := range directions {
			newRow := current.row + dir[0]
			newCol := current.col + dir[1]

			if newRow < 0 ||
				newRow >= rows ||
				newCol < 0 ||
				newCol >= cols ||
				visited[newRow][newCol] {
				continue
			}

			visited[newRow][newCol] = true

			neighborHeight := heightMap[newRow][newCol]

			if neighborHeight < current.height {
				totalWater += current.height - neighborHeight
			}

			effectiveHeight := max(current.height, neighborHeight)

			heap.Push(
				minHeap,
				Cell{effectiveHeight, newRow, newCol},
			)
		}
	}

	return totalWater
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

The Go solution mirrors the Python logic closely, but Go requires a custom heap implementation using the container/heap package.

Unlike Python's built in heapq, Go requires defining:

  • Len
  • Less
  • Swap
  • Push
  • Pop

The visited matrix is implemented using slices of boolean slices.

Go integers are sufficient for this problem because the maximum trapped water remains well within 32 bit integer range.

Worked Examples

Example 1

Input:

[
  [1,4,3,1,3,2],
  [3,2,1,3,2,4],
  [2,3,3,2,3,1]
]

Initial border heap:

Height Position
1 (0,0)
1 (0,3)
1 (2,5)
2 (0,5)
2 (2,0)
3 ...

The heap always processes the smallest boundary first.

Step 1

Pop:

(1, 0, 0)

Neighbors are already borders or too high.

No water trapped.

Step 2

Pop:

(1, 0, 3)

Explore neighbor:

(1, 3)

Height is 3, no water trapped.

Explore:

(1, 2)

Height is 1.

Current boundary height is 1.

Water trapped:

1 - 1 = 0

Push effective height:

max(1, 1) = 1

Step 3

Pop:

(1, 1, 2)

Neighbor:

(1,1)

Height = 2

Water trapped:

max(0, 1 - 2) = 0

Neighbor:

(1,4)

Height = 2

Still no water.

As the algorithm continues, lower interior cells surrounded by taller boundaries trap water.

Final result:

4

Example 2

Input:

[
  [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 outer boundary is entirely height 3.

The algorithm gradually expands inward.

Interior Processing

Cell Height Boundary Water Trapped
(1,1) 2 3 1
(1,2) 2 3 1
(1,3) 2 3 1
(2,1) 2 3 1
(2,2) 1 3 2
(2,3) 2 3 1
(3,1) 2 3 1
(3,2) 2 3 1
(3,3) 2 3 1

Total:

1+1+1+1+2+1+1+1+1 = 10

Final result:

10

Complexity Analysis

Measure Complexity Explanation
Time O(m × n × log(m × n)) Each cell is inserted and removed from the heap once
Space O(m × n) Visited matrix and heap storage

The heap contains at most m × n cells. Every push and pop operation costs logarithmic time.

Since each cell is processed once, the total number of heap operations is proportional to the number of cells.

Therefore the final complexity is:

$$O(m \times n \times \log(m \times n))$$

which is efficient enough for grids up to 200 x 200.

Test Cases

from typing import List

class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        import heapq

        if not heightMap or not heightMap[0]:
            return 0

        rows = len(heightMap)
        cols = len(heightMap[0])

        if rows < 3 or cols < 3:
            return 0

        visited = [[False] * cols for _ in range(rows)]
        heap = []

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

        for c in range(cols):
            for r in [0, rows - 1]:
                if not visited[r][c]:
                    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:
            h, r, c = heapq.heappop(heap)

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

                if (
                    nr < 0 or nr >= rows or
                    nc < 0 or nc >= cols or
                    visited[nr][nc]
                ):
                    continue

                visited[nr][nc] = True

                nh = heightMap[nr][nc]

                if nh < h:
                    water += h - nh

                heapq.heappush(heap, (max(h, nh), nr, nc))

        return water

sol = Solution()

assert sol.trapRainWater(
    [[1,4,3,1,3,2],
     [3,2,1,3,2,4],
     [2,3,3,2,3,1]]
) == 4  # Provided example 1

assert sol.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  # Provided example 2

assert sol.trapRainWater([[1]]) == 0  # Single cell

assert sol.trapRainWater([[1,2,3,4]]) == 0  # Single row

assert sol.trapRainWater([[1],[2],[3]]) == 0  # Single column

assert sol.trapRainWater(
    [[5,5,5],
     [5,1,5],
     [5,5,5]]
) == 4  # Simple basin

assert sol.trapRainWater(
    [[5,5,5],
     [5,5,5],
     [5,5,5]]
) == 0  # Flat terrain

assert sol.trapRainWater(
    [[5,5,5,5],
     [5,1,2,5],
     [5,5,5,5]]
) == 7  # Multiple trapped cells

assert sol.trapRainWater(
    [[12,13,1,12],
     [13,4,13,12],
     [13,8,10,12],
     [12,13,12,12],
     [13,13,13,13]]
) == 14  # Classic tricky case

assert sol.trapRainWater(
    [[9,9,9,9],
     [9,0,0,9],
     [9,0,0,9],
     [9,9,9,9]]
) == 36  # Large enclosed basin

Test Case Summary

Test Why
Example 1 Validates standard mixed terrain
Example 2 Validates large enclosed basin
Single cell No interior region exists
Single row Water cannot be enclosed vertically
Single column Water cannot be enclosed horizontally
Simple basin Basic trapped water scenario
Flat terrain Ensures no false positives
Multiple trapped cells Tests accumulation logic
Classic tricky case Validates global boundary handling
Large enclosed basin Tests large water volume accumulation

Edge Cases

Very Small Grids

Grids with fewer than three rows or columns cannot trap water because there is no enclosed interior space.

For example:

[[1,2,3]]

Every cell lies on the boundary, so water always escapes immediately.

The implementation explicitly checks:

if rows < 3 or cols < 3:
    return 0

This avoids unnecessary heap processing.

Flat Terrain

A completely flat map traps no water because no cell is lower than its surrounding boundary.

Example:

[
  [5,5,5],
  [5,5,5],
  [5,5,5]
]

A buggy implementation might incorrectly assume the center can hold water simply because it is enclosed.

The heap approach avoids this because trapped water is only added when:

neighbor_height < current_height

Equal heights trap nothing.

Hidden Escape Paths

This is the most subtle case.

Consider a region that appears enclosed locally but has a distant low boundary somewhere else.

A naive local approach might overestimate trapped water because it only looks at nearby walls.

The heap based BFS solves this correctly because expansion always starts from the globally smallest boundary. If water can escape through a lower route anywhere on the map, that route will be processed before higher boundaries.

This guarantees that the algorithm never traps more water than physically possible.