Want Help Cracking FAANG?

(Then click this)

 ×
Back to Question Bank

407. Trapping Rain Water II - Leetcode Solution

Code Implementation

import heapq

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

        m, n = len(heightMap), len(heightMap[0])
        visited = [[False] * n for _ in range(m)]
        heap = []
        
        # Add all the boundary cells to the heap
        for i in range(m):
            for j in [0, n-1]:
                heapq.heappush(heap, (heightMap[i][j], i, j))
                visited[i][j] = True
        for j in range(n):
            for i in [0, m-1]:
                if not visited[i][j]:
                    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:
            height, x, y = heapq.heappop(heap)
            for dx, dy in dirs:
                nx, ny = x + dx, y + dy
                if 0 <= nx < m and 0 <= ny < n and not visited[nx][ny]:
                    visited[nx][ny] = True
                    water += max(0, height - heightMap[nx][ny])
                    heapq.heappush(heap, (max(height, heightMap[nx][ny]), nx, ny))
        return water
      
#include <vector>
#include <queue>
using namespace std;

class Solution {
public:
    int trapRainWater(vector<vector<int>>& heightMap) {
        int m = heightMap.size();
        if (m == 0) return 0;
        int n = heightMap[0].size();
        if (n == 0) return 0;
        vector<vector<bool>> visited(m, vector<bool>(n, false));
        using Cell = tuple<int, int, int>; // height, x, y
        priority_queue<Cell, vector<Cell>, greater<Cell>> heap;

        // Push boundary cells
        for (int i = 0; i < m; ++i) {
            heap.emplace(heightMap[i][0], i, 0);
            visited[i][0] = true;
            heap.emplace(heightMap[i][n-1], i, n-1);
            visited[i][n-1] = true;
        }
        for (int j = 1; j < n-1; ++j) {
            heap.emplace(heightMap[0][j], 0, j);
            visited[0][j] = true;
            heap.emplace(heightMap[m-1][j], m-1, j);
            visited[m-1][j] = true;
        }

        int dirs[4][2] = {{-1,0}, {1,0}, {0,-1}, {0,1}};
        int water = 0;
        while (!heap.empty()) {
            auto [h, x, y] = heap.top(); heap.pop();
            for (auto& d : dirs) {
                int nx = x + d[0], ny = y + d[1];
                if (nx < 0 || nx >= m || ny < 0 || ny >= n || visited[nx][ny]) continue;
                visited[nx][ny] = true;
                water += max(0, h - heightMap[nx][ny]);
                heap.emplace(max(h, heightMap[nx][ny]), nx, ny);
            }
        }
        return water;
    }
};
      
import java.util.*;

class Solution {
    public int trapRainWater(int[][] heightMap) {
        int m = heightMap.length;
        if (m == 0) return 0;
        int n = heightMap[0].length;
        if (n == 0) return 0;

        boolean[][] visited = new boolean[m][n];
        PriorityQueue<int[]> heap = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));

        // Push boundary cells
        for (int i = 0; i < m; i++) {
            heap.offer(new int[]{heightMap[i][0], i, 0});
            visited[i][0] = true;
            heap.offer(new int[]{heightMap[i][n-1], i, n-1});
            visited[i][n-1] = true;
        }
        for (int j = 1; j < n-1; j++) {
            heap.offer(new int[]{heightMap[0][j], 0, j});
            visited[0][j] = true;
            heap.offer(new int[]{heightMap[m-1][j], m-1, j});
            visited[m-1][j] = true;
        }

        int[][] dirs = {{-1,0}, {1,0}, {0,-1}, {0,1}};
        int water = 0;
        while (!heap.isEmpty()) {
            int[] cell = heap.poll();
            int h = cell[0], x = cell[1], y = cell[2];
            for (int[] d : dirs) {
                int nx = x + d[0], ny = y + d[1];
                if (nx < 0 || nx >= m || ny < 0 || ny >= n || visited[nx][ny]) continue;
                visited[nx][ny] = true;
                water += Math.max(0, h - heightMap[nx][ny]);
                heap.offer(new int[]{Math.max(h, heightMap[nx][ny]), nx, ny});
            }
        }
        return water;
    }
}
      
class MinHeap {
    constructor() { this.data = []; }
    push(val) {
        this.data.push(val);
        this._bubbleUp(this.data.length - 1);
    }
    pop() {
        if (this.data.length === 1) return this.data.pop();
        const top = this.data[0];
        this.data[0] = this.data.pop();
        this._bubbleDown(0);
        return top;
    }
    _bubbleUp(idx) {
        while (idx > 0) {
            let parent = Math.floor((idx - 1) / 2);
            if (this.data[parent][0] <= this.data[idx][0]) break;
            [this.data[parent], this.data[idx]] = [this.data[idx], this.data[parent]];
            idx = parent;
        }
    }
    _bubbleDown(idx) {
        let length = this.data.length;
        while (true) {
            let left = idx * 2 + 1, right = idx * 2 + 2, smallest = idx;
            if (left < length && this.data[left][0] < this.data[smallest][0]) smallest = left;
            if (right < length && this.data[right][0] < this.data[smallest][0]) smallest = right;
            if (smallest === idx) break;
            [this.data[smallest], this.data[idx]] = [this.data[idx], this.data[smallest]];
            idx = smallest;
        }
    }
    isEmpty() { return this.data.length === 0; }
}

var trapRainWater = function(heightMap) {
    if (!heightMap.length || !heightMap[0].length) return 0;
    let m = heightMap.length, n = heightMap[0].length;
    let visited = Array.from({length: m}, () => Array(n).fill(false));
    let heap = new MinHeap();

    // Push boundary cells
    for (let i = 0; i < m; i++) {
        heap.push([heightMap[i][0], i, 0]);
        visited[i][0] = true;
        heap.push([heightMap[i][n-1], i, n-1]);
        visited[i][n-1] = true;
    }
    for (let j = 1; j < n-1; j++) {
        heap.push([heightMap[0][j], 0, j]);
        visited[0][j] = true;
        heap.push([heightMap[m-1][j], m-1, j]);
        visited[m-1][j] = true;
    }

    let dirs = [[-1,0],[1,0],[0,-1],[0,1]];
    let water = 0;
    while (!heap.isEmpty()) {
        let [h, x, y] = heap.pop();
        for (let [dx, dy] of dirs) {
            let nx = x + dx, ny = y + dy;
            if (nx < 0 || nx >= m || ny < 0 || ny >= n || visited[nx][ny]) continue;
            visited[nx][ny] = true;
            water += Math.max(0, h - heightMap[nx][ny]);
            heap.push([Math.max(h, heightMap[nx][ny]), nx, ny]);
        }
    }
    return water;
};

Problem Description

Given a 2D integer matrix heightMap representing an elevation map where heightMap[i][j] is the height of the terrain at cell (i, j), compute how much water it is able to trap after raining.
  • The water can only be trapped in "basins"—areas surrounded on all sides by higher elevation.
  • Water cannot leak through the edges of the map.
  • Each cell is surrounded by four neighbors (up, down, left, right).
  • Return the total units of water that can be trapped.
Constraints:
  • m == heightMap.length
  • n == heightMap[i].length
  • 1 <= m, n <= 200
  • 0 <= heightMap[i][j] <= 2 * 10^4

Thought Process

When first approaching this problem, you might think of it as an extension of the 1D "Trapping Rain Water" problem. In 1D, you can precompute the left and right maximums for each cell and use them to calculate trapped water. However, in 2D, water can escape in any of four directions, making it much more complex. A brute-force approach would be to, for each cell, find the minimum boundary height in all directions, but this is extremely inefficient for large grids. Instead, we need a way to efficiently process the boundaries and propagate the minimum "wall" height inward, similar to how water would actually fill a basin. The key insight is that water will always be trapped by the lowest boundary around it. By processing the lowest boundary cells first, and moving inward, we can determine for each cell the maximum water it can hold, based on the minimum height of its surrounding walls.

Solution Approach

To solve this problem efficiently, we use a priority queue (min-heap) to always process the lowest boundary first, and a visited matrix to keep track of processed cells. Here’s the step-by-step approach:
  1. Initialization:
    • Add all the boundary cells (edges of the matrix) into a min-heap. Mark them as visited.
    • The heap stores tuples: (height, x, y).
  2. Processing:
    • While the heap is not empty, pop the cell with the lowest height.
    • For each of its four neighbors (up, down, left, right):
      • If the neighbor has not been visited:
        • The water trapped at this neighbor is max(0, current_height - neighbor_height).
        • Add this water to the result.
        • Push the neighbor into the heap with max(current_height, neighbor_height) as its new height (the wall is now at least as high as the highest of the two).
        • Mark the neighbor as visited.
  3. Repeat until all cells are visited.
Why this works:
  • By always expanding from the lowest boundary, we ensure that we never overestimate the amount of water that can be trapped at any cell.
  • Each cell is processed only once, and the heap ensures we always process the next "lowest wall" first, mimicking how water would actually fill up basins.

Example Walkthrough

Input: 
  heightMap = [
    [1, 4, 3, 1, 3, 2],
    [3, 2, 1, 3, 2, 4],
    [2, 3, 3, 2, 3, 1]
  ]
Step-by-step:
  1. Push all boundary cells into the heap:
    • First and last rows: all cells
    • First and last columns of middle rows
  2. Start popping the lowest cell from the heap, and check its neighbors.
  3. For cell (1,2) with height 1, surrounded by higher boundaries, water can be trapped:
    • It is surrounded by boundaries of height 3 and 2, so it can trap min(3,2) - 1 = 1 unit of water.
  4. Continue this process for all cells. The total trapped water is 4 units for this example.
Why? At each step, you only trap water if the current cell is lower than the minimum boundary height encountered so far.

Time and Space Complexity

  • Brute-force approach: For each cell, you would have to scan outwards to find the minimum boundary height, which is O((mn) * (m + n)) and infeasible for large grids.
  • Optimized (heap-based) approach:
    • Each cell is pushed and popped from the heap at most once: O(mn log(mn)) due to heap operations.
    • Space: O(mn) for the visited matrix and the heap.
Why? Heap operations are logarithmic in the number of elements, and each element is processed once.

Summary

The "Trapping Rain Water II" problem generalizes the classic 1D trapping water problem to two dimensions. By leveraging a priority queue to always process the lowest boundary cell and expanding inward, we can efficiently and correctly compute the amount of water trapped in complex 2D basins. The approach is both elegant and optimal for large inputs, combining ideas from BFS and greedy algorithms, and avoids the inefficiency of brute-force methods.