Trapping Rain Water II — 3D Min-Heap BFS Interview Pattern

Problem Statement

Given an m x n integer matrix heightMap, return how much water can be trapped after raining. Water cannot stay on border cells; only inner cells trap if surrounded by taller cells.

Constraints:

m == heightMap.length, n == heightMap[0].length
1 <= m, n <= 200
0 <= heightMap[i][j] <= 2 * 10^4

Input:  [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]
Output: 4

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]]
Output: 10

Why This Problem Matters

LeetCode 407 is the 3D extension of the classic Trapping Rain Water problem and shows up in Google, Meta, and Amazon hard interview rounds. It is the textbook example of using a min-heap to drive a BFS frontier — a pattern reused in shortest-path-like grid problems.

The "process the lowest boundary first" idea is one of the most elegant uses of a priority queue interview tool. Once you internalize it, you will recognize it instantly in problems like Swim in Rising Water, Path with Minimum Effort, and watershed segmentation in image processing.

The Core Insight

Water trapped at any cell equals the lowest border height that surrounds it minus the cell height. Push all border cells into a min-heap and BFS inward. For each popped cell, water at each unvisited neighbor is max(0, popped_height - neighbor_height). Push the neighbor with max(popped_height, neighbor_height) because that is the new effective wall height.

Visual Dry Run

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

Step	Pop (h, r, c)	Neighbor	Trap	Push
1	(1, 0, 0)	(1, 0): 4	0	(4, 0, 1)
2	(1, 0, 3)	(1, 3): 3	0	(3, 1, 3)
3	(1, 2, 5)	(1, 5): 4	0	(4, 1, 5)
4	(2, 0, 5)	(1, 1): 2	1	wall=2
5	(2, 1, 4)	(1, 4): 2	1	wall=2
6	...	...	...	...
Total			4

Solution (Optimal)

import heapq
from typing import List
 
class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        if not heightMap or len(heightMap) < 3 or len(heightMap[0]) < 3:
            return 0
        m, n = len(heightMap), len(heightMap[0])
        visited = [[False] * n for _ in range(m)]
        heap = []
        for i in range(m):
            for j in range(n):
                if i == 0 or i == m - 1 or j == 0 or j == n - 1:
                    heapq.heappush(heap, (heightMap[i][j], i, j))
                    visited[i][j] = True
        water = 0
        dirs = [(-1, 0), (1, 0), (0, -1), (0, 1)]
        while heap:
            h, r, c = heapq.heappop(heap)
            for dr, dc in dirs:
                nr, nc = r + dr, c + dc
                if 0 <= nr < m and 0 <= nc < n and not visited[nr][nc]:
                    visited[nr][nc] = True
                    water += max(0, h - heightMap[nr][nc])
                    heapq.heappush(heap, (max(h, heightMap[nr][nc]), nr, nc))
        return water

class MinHeap {
    constructor() { this.h = []; }
    push(v) { this.h.push(v); this._up(this.h.length - 1); }
    pop() {
        const top = this.h[0], last = this.h.pop();
        if (this.h.length) { this.h[0] = last; this._down(0); }
        return top;
    }
    _up(i) {
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (this.h[i][0] < this.h[p][0]) {
                [this.h[i], this.h[p]] = [this.h[p], this.h[i]];
                i = p;
            } else break;
        }
    }
    _down(i) {
        const n = this.h.length;
        while (true) {
            const l = 2 * i + 1, r = 2 * i + 2;
            let m = i;
            if (l < n && this.h[l][0] < this.h[m][0]) m = l;
            if (r < n && this.h[r][0] < this.h[m][0]) m = r;
            if (m === i) break;
            [this.h[i], this.h[m]] = [this.h[m], this.h[i]];
            i = m;
        }
    }
    get size() { return this.h.length; }
}
 
var trapRainWater = function(heightMap) {
    const m = heightMap.length, n = heightMap[0].length;
    if (m < 3 || n < 3) return 0;
    const visited = Array.from({ length: m }, () => new Array(n).fill(false));
    const heap = new MinHeap();
    for (let i = 0; i < m; i++) {
        for (let j = 0; j < n; j++) {
            if (i === 0 || i === m - 1 || j === 0 || j === n - 1) {
                heap.push([heightMap[i][j], i, j]);
                visited[i][j] = true;
            }
        }
    }
    let water = 0;
    const dirs = [[-1, 0], [1, 0], [0, -1], [0, 1]];
    while (heap.size) {
        const [h, r, c] = heap.pop();
        for (const [dr, dc] of dirs) {
            const nr = r + dr, nc = c + dc;
            if (nr >= 0 && nr < m && nc >= 0 && nc < n && !visited[nr][nc]) {
                visited[nr][nc] = true;
                water += Math.max(0, h - heightMap[nr][nc]);
                heap.push([Math.max(h, heightMap[nr][nc]), nr, nc]);
            }
        }
    }
    return water;
};

Time: O(mn log(mn)) — every cell pushed/popped once with log(mn) heap ops. Space: O(mn) — visited matrix plus heap.

Common Mistakes

Pushing inner cells initially — only the border can hold water boundary.
Forgetting to mark visited at push time, not pop time; otherwise a cell enters the heap multiple times.
Pushing heightMap[nr][nc] instead of max(h, height) after trapping — destroys the wall height.
Not handling grids smaller than 3x3 (no inner cells means 0 water).
Treating border cells as trappable — they leak.

Interview Tips

Connect this to the 1D Trapping Rain Water two-pointer solution and show why 2D needs a heap.
Explain why mark-on-push not pop: O(mn) cells, O(1) entries each.
Mention watershed segmentation as a real-world use case.
Discuss why DP fails here: water at a cell depends on all directions, not just left/right.

Follow-up Questions

What if the grid is enormous and cannot fit in memory? Hint: tile-based out-of-core processing.
Can you do it without max(h, height)? Hint: no — that wall height is the invariant.
How would you find which cells contribute to the leak edge? Hint: track the popping order.
What about floating-point heights? Hint: same algorithm; use a comparator with epsilon.
Find max water depth at any single cell. Hint: compute per-cell, not sum.

Key Takeaways

LeetCode 407 Trapping Rain Water II uses min-heap BFS in O(mn log(mn)).
Always start the heap with all border cells.
Push neighbors with max(current_wall, neighbor_height) — this preserves the seal.
Mark visited at push time, never at pop time.
The pattern generalizes to watershed and image segmentation algorithms.
A 1D two-pointer approach does not extend to 2D; the heap is essential.
This problem is a Google and Meta hard interview favorite for grid plus heap reasoning.

Problem Statement#

Why This Problem Matters#

The Core Insight#

Visual Dry Run#

Solution (Optimal)#

Common Mistakes#

Interview Tips#

Follow-up Questions#

Key Takeaways#

Related reading