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.

Note: Both m and n are less than 110. The height of each unit cell is greater than 0 and is less than 20,000.

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.

Solution:

public class Solution {
    public class Cell {
        int x;
        int y;
        int h;
        public Cell(int x, int y, int h){
            this.x = x;
            this.y = y;
            this.h = h;
        }
    }
    public int trapRainWater(int[][] heightMap) {
        int m = heightMap.length;
        if(m==0) return 0;
        int n = heightMap[0].length;
        boolean[][] visited = new boolean[m][n];
        PriorityQueue<Cell> pq = new PriorityQueue<>(2*(m+n), (a,b)->(a.h-b.h));
        for(int i=0; i<m; i++){
            visited[i][0] = true;
            pq.offer(new Cell(i, 0, heightMap[i][0]));
            visited[i][n-1] = true;
            pq.offer(new Cell(i, n-1, heightMap[i][n-1]));
        }
        for(int j=1; j<n-1; j++){
            visited[0][j] = true;
            pq.offer(new Cell(0, j, heightMap[0][j]));
            visited[m-1][j] = true;
            pq.offer(new Cell(m-1, j, heightMap[m-1][j]));
        }
        int ans = 0;
        int[][] dirs = new int[][]{{-1, 0}, {1, 0}, {0, -1}, {0, 1}};
        while(!pq.isEmpty()){
            Cell cell = pq.poll();
            for(int[] d : dirs){
                int row = cell.x + d[0];
                int col = cell.y + d[1];
                if(row>0 && row<m-1 && col>0 && col<n-1 && !visited[row][col]){
                    visited[row][col] = true;
                    ans += Math.max(0, cell.h - heightMap[row][col]);
                    pq.offer(new Cell(row, col, Math.max(cell.h, heightMap[row][col]))); /* maximum */
                }
            }
            
        }
        return ans;
    }
}
comments powered by Disqus