Trapping Rain Water II

Given an m x n matrix of positive integers representing the height of each unit cell in a 2D elevation map, compute the volume of water it is able to trap after raining.

Example:

Given the following 3x6 height map:
[
  [1,4,3,1,3,2],
  [3,2,1,3,2,4],
  [2,3,3,2,3,1]
]

Return 4.

The above image represents the elevation map [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]] before the rain.

After the rain, water is trapped between the blocks. The total volume of water trapped is 4.

Constraints:

  • 1 <= m, n <= 110

  • 0 <= heightMap[i][j] <= 20000

class Solution {
    // https://drive.google.com/file/d/1YHHlkHkx5Ae9Ft9dAHMV52HMnOCnYcMF/view?usp=sharing
    static class Pair {
        int x, y, val;

        Pair(int x, int y, int v) {
            this.x = x;
            this.y = y;
            this.val = v;
        }
    }

    public int trapRainWater(int[][] heightMap) {
        int m = heightMap.length, n = heightMap[0].length;
        if (m == 1 || n == 1)
            return 0;
        // Priority Queue hold virtual border which holds water inside
        // Therefore starting border is border of matrix
        PriorityQueue<Pair> pq = new PriorityQueue<>((a, b) -> a.val - b.val);
        boolean[][] visited = new boolean[m][n];
        for (int i = 0; i < m; i++) {
            pq.add(new Pair(i, 0, heightMap[i][0]));
            visited[i][0] = true;
            pq.add(new Pair(i, n - 1, heightMap[i][n - 1]));
            visited[i][n - 1] = true;
        }
        for (int i = 1; i < n - 1; i++) {
            pq.add(new Pair(0, i, heightMap[0][i]));
            visited[0][i] = true;
            pq.add(new Pair(m - 1, i, heightMap[m - 1][i]));
            visited[m - 1][i] = true;
        }
        // Now we want to find smallest height, and check all unvisited neighbors
        // If any neighbor < border height(min), then we can add this water to our
        // answer, as it cannot flow out of border from any direction
        int[][] dirs = { { 1, 0 }, { 0, 1 }, { -1, 0 }, { 0, -1 } };
        int ans = 0;
        while (pq.size() != 0) {
            Pair smallest = pq.poll();
            int x = smallest.x, y = smallest.y;
            for (int[] dir : dirs) {
                int newX = x + dir[0], newY = y + dir[1];
                if (newX >= 0 && newX < m && newY >= 0 && newY < n && !visited[newX][newY]) {
                    visited[newX][newY] = true;
                    // If we can store some water here, then this stored water
                    // will also work as a wall in our virtual border
                    if (heightMap[newX][newY] < smallest.val) {
                        ans += smallest.val - heightMap[newX][newY];
                        pq.add(new Pair(newX, newY, smallest.val));
                    } else
                        pq.add(new Pair(newX, newY, heightMap[newX][newY]));
                }
            }
        }
        return ans;
    }
}

Last updated