Trapping Rain Water II — Hard Problem & Solution

An m × n matrix gives the height of each unit cell of a 2D elevation map. Return the volume of water it can trap after raining.

Difficulty: Hard
Topics: Arrays, Matrix, Breadth-First Search, Heap
Asked at: Amazon, Google, Twitter
Time limit: 2 s
Memory limit: 256 MB
Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

An m × n matrix gives the height of each unit cell of a 2D elevation map. Return the volume of water it can trap after raining. Water runs off the edges of the map.

Example 1

Input: heightMap = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]
Output: 4
Explanation: The four interior cells hold 4 units between them.

Example 2

Input: 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]]
Output: 10
Explanation: A wall of 3 all round: eight cells fill by 1 and the centre by 2.

Example 3

Input: heightMap = [[1,1,1],[1,0,1],[1,1,1]]
Output: 1

Constraints

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

How to solve Trapping Rain Water II

Flood inward from the boundary, always processing the lowest cell on the current rim. When the rim's lowest cell has level L, any unvisited neighbour below L is filled to L, because L is the cheapest way out of the basin.

Approach

Push every border cell onto a min-heap keyed by its height, and mark it visited.
Pop the lowest cell, with water level L.
For each unvisited neighbour: if its height is below L, it traps L - height; push it back at level max(L, height).
Continue until the heap empties.

Why it works

Processing the rim in increasing order of level is the whole argument: when a cell is popped at level L, every other route out of the region is at least L high, so L is exactly the water level at its neighbours — no later discovery can lower it. Pushing a neighbour at max(L, height) carries that boundary forward, since a taller cell raises the rim rather than letting water out.

Complexity

Time — O(m · n · log(m · n))
Space — O(m · n)

Pitfalls

Marking a cell visited when it is pushed — not when popped — keeps each cell out of the heap more than once.
A grid narrower than 3 in either direction traps nothing; it is all border.
The pushed level is max(level, height), not the raw height, or a basin behind a tall wall leaks.

Reference solution

Python

from typing import List
import heapq

def trapRainWater(heightMap: List[List[int]]) -> int:
    m, n = len(heightMap), len(heightMap[0])
    if m < 3 or n < 3:
        return 0
    seen = [[False] * n for _ in range(m)]
    heap = []
    for i in range(m):
        for j in range(n):
            if i not in (0, m - 1) and j not in (0, n - 1):
                continue
            seen[i][j] = True
            heapq.heappush(heap, (heightMap[i][j], i, j))
    total = 0
    while heap:
        level, r, c = heapq.heappop(heap)
        for nr, nc in ((r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1)):
            if not (0 <= nr < m and 0 <= nc < n) or seen[nr][nc]:
                continue
            seen[nr][nc] = True
            if heightMap[nr][nc] < level:
                total += level - heightMap[nr][nc]
            heapq.heappush(heap, (max(level, heightMap[nr][nc]), nr, nc))
    return total

JavaScript

var trapRainWater = function(heightMap) {
    var m = heightMap.length, n = heightMap[0].length, i, j;
    if (m < 3 || n < 3) return 0;
    var seen = [];
    for (i = 0; i < m; i++) {
        var row = [];
        for (j = 0; j < n; j++) row.push(false);
        seen.push(row);
    }
    // The heap holds level * 40000 + cell, so a plain int heap is enough:
    // m * n <= 40000 and the level fits well inside 32 bits.
    var heap = [];
    var push = function(v) {
        heap.push(v);
        var k = heap.length - 1;
        while (k > 0) {
            var p = (k - 1) >> 1;
            if (heap[p] <= heap[k]) break;
            var t = heap[p]; heap[p] = heap[k]; heap[k] = t;
            k = p;
        }
    };
    var pop = function() {
        var top = heap[0];
        var last = heap.pop();
        if (heap.length > 0) {
            heap[0] = last;
            var k = 0;
            for (;;) {
                var l = 2 * k + 1, r = l + 1, s = k;
                if (l < heap.length && heap[l] < heap[s]) s = l;
                if (r < heap.length && heap[r] < heap[s]) s = r;
                if (s === k) break;
                var t = heap[s]; heap[s] = heap[k]; heap[k] = t;
                k = s;
            }
        }
        return top;
    };
    for (i = 0; i < m; i++) {
        for (j = 0; j < n; j++) {
            if (i !== 0 && i !== m - 1 && j !== 0 && j !== n - 1) continue;
            seen[i][j] = true;
            push(heightMap[i][j] * 40000 + i * n + j);
        }
    }
    var dr = [1, -1, 0, 0], dc = [0, 0, 1, -1];
    var total = 0;
    while (heap.length > 0) {
        var key = pop();
        var level = Math.floor(key / 40000);
        var cell = key % 40000;
        var cr = Math.floor(cell / n), cc = cell % n;
        for (var d = 0; d < 4; d++) {
            var nr = cr + dr[d], nc = cc + dc[d];
            if (nr < 0 || nr >= m || nc < 0 || nc >= n || seen[nr][nc]) continue;
            seen[nr][nc] = true;
            if (heightMap[nr][nc] < level) total += level - heightMap[nr][nc];
            var next = heightMap[nr][nc] > level ? heightMap[nr][nc] : level;
            push(next * 40000 + nr * n + nc);
        }
    }
    return total;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 667 arrays problems · the whole catalogue