Trapping Rain Water II

hard heap bfs matrix

Problem

Given an m×n elevation map, compute the volume of water it can trap after raining.

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

Output4

A min-heap of border cells lets the algorithm flood inward in non-decreasing barrier order.

heights (rows ';' cells ',')

import heapq
def trap_rain_water(h):
    m, n = len(h), len(h[0])
    if m < 3 or n < 3: return 0
    heap = []
    seen = [[False]*n for _ in range(m)]
    for r in range(m):
        for c in range(n):
            if r in (0, m-1) or c in (0, n-1):
                heapq.heappush(heap, (h[r][c], r, c))
                seen[r][c] = True
    water = 0
    while heap:
        ht, r, c = heapq.heappop(heap)
        for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
            nr, nc = r + dr, c + dc
            if 0 <= nr < m and 0 <= nc < n and not seen[nr][nc]:
                seen[nr][nc] = True
                water += max(0, ht - h[nr][nc])
                heapq.heappush(heap, (max(ht, h[nr][nc]), nr, nc))
    return water

function trapRainWater(h) {
  const m = h.length, n = h[0].length;
  if (m < 3 || n < 3) return 0;
  const heap = [];
  const seen = Array.from({length: m}, () => new Array(n).fill(false));
  function push(p) { heap.push(p); heap.sort((a, b) => a[0] - b[0]); }
  function pop() { return heap.shift(); }
  for (let r = 0; r < m; r++) for (let c = 0; c < n; c++)
    if (r === 0 || r === m - 1 || c === 0 || c === n - 1) { push([h[r][c], r, c]); seen[r][c] = true; }
  let water = 0;
  while (heap.length) {
    const [ht, r, c] = pop();
    for (const [dr, dc] of [[-1,0],[1,0],[0,-1],[0,1]]) {
      const nr = r + dr, nc = c + dc;
      if (nr < 0 || nr >= m || nc < 0 || nc >= n || seen[nr][nc]) continue;
      seen[nr][nc] = true;
      water += Math.max(0, ht - h[nr][nc]);
      push([Math.max(ht, h[nr][nc]), nr, nc]);
    }
  }
  return water;
}

class Solution {
    public int trapRainWater(int[][] h) {
        int m = h.length, n = h[0].length;
        if (m < 3 || n < 3) return 0;
        PriorityQueue heap = new PriorityQueue<>((a, b) -> a[0] - b[0]);
        boolean[][] seen = new boolean[m][n];
        for (int r = 0; r < m; r++) for (int c = 0; c < n; c++)
            if (r == 0 || r == m - 1 || c == 0 || c == n - 1) { heap.add(new int[]{ h[r][c], r, c }); seen[r][c] = true; }
        int[][] D = {{-1,0},{1,0},{0,-1},{0,1}};
        int water = 0;
        while (!heap.isEmpty()) {
            int[] cur = heap.poll();
            for (int[] d : D) {
                int nr = cur[1] + d[0], nc = cur[2] + d[1];
                if (nr < 0 || nr >= m || nc < 0 || nc >= n || seen[nr][nc]) continue;
                seen[nr][nc] = true;
                water += Math.max(0, cur[0] - h[nr][nc]);
                heap.add(new int[]{ Math.max(cur[0], h[nr][nc]), nr, nc });
            }
        }
        return water;
    }
}

int trapRainWater(vector>& h) {
    int m = h.size(), n = h[0].size();
    if (m < 3 || n < 3) return 0;
    priority_queue, vector>, greater<>> heap;
    vector> seen(m, vector(n, false));
    for (int r = 0; r < m; r++) for (int c = 0; c < n; c++)
        if (r == 0 || r == m - 1 || c == 0 || c == n - 1) { heap.push({h[r][c], r, c}); seen[r][c] = true; }
    int dr[] = {-1,1,0,0}, dc[] = {0,0,-1,1};
    int water = 0;
    while (!heap.empty()) {
        auto [ht, r, c] = heap.top(); heap.pop();
        for (int k = 0; k < 4; k++) {
            int nr = r + dr[k], nc = c + dc[k];
            if (nr < 0 || nr >= m || nc < 0 || nc >= n || seen[nr][nc]) continue;
            seen[nr][nc] = true;
            water += max(0, ht - h[nr][nc]);
            heap.push({max(ht, h[nr][nc]), nr, nc});
        }
    }
    return water;
}

Explanation

Water on a 2D map is held in by the lowest wall around it. The trick is to start from the outer border cells (which can never hold water) and flood inward, always processing the shortest wall first using a min-heap.

We push every border cell into a min-heap keyed by height and mark them seen. We then repeatedly pop the cell with the smallest height. That popped height acts as the current water barrier for its unvisited neighbors.

For each new neighbor, if its own height is below the barrier ht, the difference ht - h[nr][nc] is trapped water. Either way we push the neighbor back with height max(ht, h[nr][nc]), because a low cell still gets walled in at the barrier level.

Processing the smallest wall first guarantees that when we reach an inner cell, we already know the lowest surrounding rim, so the water amount is correct.

Example: in the sample grid the inner dips collect a total of 4 units, each found as a neighbor sits below the smallest enclosing border height.

Time: O(mn log(mn)) Space: O(mn)