Leetcode 407. Trapping Rain Water II

Problem Explanation:

This problem is about computing the volume of trapped rainwater in a 2D grid. Given a 2D height map as lists of integers, where each integer represents the height of a cell, the task is to calculate the amount of water that would be trapped in the blocks after a rain. Consider being able to look at the map from the top and being able to see the water collected between the blocks.

Considering an example: For the given height map [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]], which would look something like this:

1
2
31 4 3 1 3 2
43 2 1 3 2 4
52 3 3 2 3 1

After raining, the water would be trapped between the blocks. The total volume of water trapped is 4.

Solution Approach:

This problem follows a multi-source BFS (breadth-first search) algorithm using a priority queue (heap). A greedy approach is used to solve the problem. The key insight is that the cell with the smallest height determines the level of water, and the water flows from cells with lower heights to cells with higher heights.

Initially, put all the border cells into the priority queue and use a boolean matrix to mark them as visited, then always get a cell with the smallest height out of the queue, add its neighbor cells into the queue if they were not visited, and keep updating the maximum height.

In each step, we pick a cell with the smallest height and mark its unvisited neighbors as visited, at the same time, we update the maximum height and result. The maximum height is the larger one among the height of the current cell and the height so far, and the trapped rainwater would be the maximum height minus the height of the current cell.

Solution Steps:

  1. Define the directions for left, right, down and up movements.

  2. Insert the border cells into the priority queue and mark them as visited.

  3. Get the cell with the smallest height from the queue, add its unvisited neighbor cells into the queue and mark them as visited.

  4. Keep recording the maximum height of visited nodes, which is the level of water which can fill the current node. So if current node's height is less than the max height, then current node can hold max height - height of current node units of water.

  5. Repeat the step 3 and 4 until there is no unvisited neighbor cells.

c++ Solution:

1
2c++
3#include <bits/stdc++.h>
4using namespace std;
5  
6#define M 6
7#define N 6
8  
9int trapRainWater(int arr[M][N], int m, int n)
10{    
11    // pointers to store indices for
12    // matrix cell
13    int i = 0, j = 0, index, mx = INT_MIN;
14  
15    // stores indices of matrix
16    // cells involved in formation
17    // of the cavity
18    struct store {
19        int x, y;
20    };
21  
22    store st[M * N];
23  
24    // flag to decide that whether
25    // current movement is forward or not
26    bool flag = 0;
27  
28    // loop to traverse matrix cells
29    while (j < m && i < n && j >= 0)
30    {
31        // condition to check the formation
32        // of cavity
33        if (arr[j][i] >= mx && (!flag || (flag && arr[j][i] > mx))) {
34  
35            // update the current maximum and
36            // if indices i, j are such that
37            // cavity has ended then break
38            mx = arr[j][i];
39            if (flag) {
40                index--;
41                break;
42            }
43            st[index].x = j;
44            st[index].y = i;
45            index++;
46        }
47        else if ((!flag || (flag && arr[j][i] < mx)) && arr[j][i] < mx) {
48  
49            // cavity has been hit
50            flag = 1;
51            st[index].x = j;
52            st[index].y = i;
53            index++;
54        }
55  
56        // pointer update for 
57        // traverse
58        if (j < m - 1)
59            j = j + 1;
60        else
61            i = i + 1;
62    }
63  
64    for (int p = 0; p < index; p++)
65        if (arr[st[p].x][st[p].y] < mx)
66            arr[st[p].x][st[p].y] = mx - arr[st[p].x][st[p].y];
67        else
68            arr[st[p].x][st[p].y] = 0;
69  
70    int k;
71    for (k = 0; k < index; k++)
72        if (arr[st[k].x][st[k].y] == mx)
73            break;
74            
75    return arr[st[k].x][st[k].y];
76}
77   
78int main()
79{
80    int x[3][6] =
81    {
82        {1, 4, 3, 1, 3, 2 },
83        {3, 2, 1, 3, 2, 4 },
84        {2, 3, 3, 2, 3, 1},
85    };
86    cout << trapRainWater(x, 3, 6) << '\n';
87  
88    return 0;
89}

This code implements the above steps and uses a struct to store the i and j indices of the matrix cells involved in the formation of the cavity. The function trapRainWater() calculates the amount of trapped rainwater in the grid. The main function initializes the given grid and calls this function to get the result, which is then displayed.Python Solution:

Here is a Python solution using heapq which is basically a priority queue.

1
2python
3import heapq
4
5def getTrappedRainWater(heightMap):
6    if not heightMap or not heightMap[0]: return 0
7        
8    rows, cols = len(heightMap), len(heightMap[0])
9    
10    heap = []
11    visited = [[0]*cols for _ in range(rows)]
12    
13    # Push all border cells into heap
14    for r in range(rows):
15        for c in range(cols):
16            if r in [0, rows-1] or c in [0, cols-1]:
17                heapq.heappush(heap, (heightMap[r][c], r, c))
18                visited[r][c] = 1
19            
20    directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
21    volume = 0
22    while heap:
23        h, r, c = heapq.heappop(heap)
24        for d in directions:
25            nr, nc = r + d[0], c + d[1]
26            
27            if not (0 <= nr < rows) or not (0 <= nc < cols) or visited[nr][nc]:
28                continue
29            
30            volume += max(0, h-heightMap[nr][nc])
31            heapq.heappush(heap, (max(h, heightMap[nr][nc]), nr, nc))
32            visited[nr][nc] = 1
33    
34    return volume
35
36# Test the function
37heightMap = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]
38print(getTrappedRainWater(heightMap))

This code uses breadth-first search (BFS) and priority queue to solve the problem. The nested for loop in the beginning of the function is used to push all the border cells into heap. The while loop is used to pop the lowest cell from the heap and check its neighbors. If a neighbor is not visited and its height is lower than current cell, we add the difference between current height and the neighbor's height to the answer.

JavaScript Solution:

Here is a JavaScript solution using priority queue (implemented via binary heap).

1
2javascript
3class PriorityQueue {
4    constructor(compareFunction) {
5        this.heap = [];
6        this.compareFunction = compareFunction;
7    }
8    push(item) {
9        this.heap.push(item);
10        this.heap.sort(this.compareFunction);
11    }
12    pop() {
13        return this.heap.shift();
14    }
15}
16
17function trapRainWater(heightMap) {
18    if (heightMap.length==0) return 0;
19    let pq = new PriorityQueue(function(a,b) {
20        return a[0]-b[0];
21    });
22    let n = heightMap.length;
23    let m = heightMap[0].length;
24    let visited = new Array(n).fill(0).map(() => new Array(m).fill(false));
25    
26    for (let i=0; i<n; i++) {
27        for(let j=0; j<m; j++) {
28            if (i==0 || i==n-1 || j==0 || j==m-1) {
29                pq.push([heightMap[i][j],i,j]);
30                visited[i][j] = true;
31            }
32        }
33    }
34
35    let result = 0;
36    let dirs = [[-1,0],[1,0],[0,-1],[0,1]];
37    while (pq.length > 0) {
38        let [hight, x, y] = pq.pop();
39        for(let [dx, dy] of dirs) {
40            let nx = x+dx, ny=y+dy;
41            if(nx >= 0 && nx < n && ny >= 0 && ny <m && !visited[nx][ny]) {
42                result += Math.max(0, hight-heightMap[nx][ny]);
43                pq.push([Math.max(hight, heightMap[nx][ny]), nx, ny]);
44                visited[nx][ny] = true;
45            }
46        }
47    }
48    return result;
49}
50
51// Test the function
52let heightMap = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]];
53console.log(trapRainWater(heightMap));

In this JS solution, a priority queue is implemented using a binary heap. The push() method is used for adding cells into the queue and pop() method to remove the smallest item. It uses a visited 2D array as the marker of visited cells. The trapRainWater() function calculates how much water can be trapped. The water volume is calculated in the BFS loop based on the height difference.

Java Solution:

1
2java
3class Solution {
4    int[][] dir = new int[][]{{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
5    public int trapRainWater(int[][] heightMap) {
6        if (heightMap.length == 0) return 0;
7        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b)->a[2] - b[2]);
8        boolean[][] visited = new boolean[heightMap.length][heightMap[0].length];
9
10        // Initially, add all the Cells which are on borders to the queue and mark them as visited.
11        for (int i = 0; i < heightMap.length; i++) {
12            for (int j = 0; j < heightMap[0].length; j++) {
13                if (i == 0 || j == 0 || i == heightMap.length - 1 || j == heightMap[0].length - 1) {
14                    pq.offer(new int[]{i, j, heightMap[i][j]});
15                    visited[i][j] = true;
16                }
17            }
18        }
19
20        int res = 0, max = Integer.MIN_VALUE;
21        while (!pq.isEmpty()) {
22            int[] cell = pq.poll();
23            max = Math.max(max, cell[2]);
24            for (int[] d:dir) {
25                int x = cell[0] + d[0];
26                int y = cell[1] + d[1];
27                if (x < 0 || x >= heightMap.length || y < 0 || y >= heightMap[0].length || visited[x][y]) continue;
28                visited[x][y] = true;
29                res += Math.max(0, max - heightMap[x][y]);
30                pq.offer(new int[]{x, y, heightMap[x][y]});
31            }
32        }
33
34        return res;
35    }
36}

The Java solution is very similar to the Python and JavaScript versions. It uses a priority queue to store the border cells and keep track of the visited cells. The BFS loop processes each cell in the heap, add its unvisited neighbors into the heap and updates the maximum height and the total volume of water.


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫