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.

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 are trapped between the blocks. The total volume of water trapped is 4.

Solution

public class Solution {

    public class Grid {
        public int x;
        public int y;
        public int h;
        public Grid(int x, int y, int h) {
            this.x = x;
            this.y = y;
            this.h = h;
        }
    }

    public int trapRainWater(int[][] heightMap) {
        if(heightMap.length <= 2 || heightMap[0].length <= 2) return 0;

        int m = heightMap.length;
        int n = heightMap[0].length;

        Queue<Grid> queue = new PriorityQueue<Grid>(new Comparator<Grid>() {
            public int compare(Grid g1, Grid g2) {
                return Integer.compare(g1.h, g2.h);
            }
        });

        boolean visited[][] = new boolean[m][n];
        for(int i = 0; i < m; i++) {
            queue.offer(new Grid(i, 0, heightMap[i][0]));
            queue.offer(new Grid(i, n - 1, heightMap[i][n-1]));
            visited[i][0] = true;
            visited[i][n-1] = true;
        }

        for(int j = 0; j < n; j++) {
            queue.offer(new Grid(0, j, heightMap[0][j]));
            queue.offer(new Grid(m-1, j, heightMap[m-1][j]));
            visited[0][j] = true;
            visited[m-1][j] = true;
        }

        int water = 0;
        int[][] dirs = new int[][]{{0,1},{0,-1},{1,0},{-1,0}};
        while(!queue.isEmpty()) {
            Grid g = queue.poll();

            // check neighbors
            for(int[] dir: dirs) {
                int x = g.x + dir[0];
                int y = g.y + dir[1];
                if(x >= 0 && x < m && y >= 0 && y < n && !visited[x][y]) {
                    int gap = g.h > heightMap[x][y] ? g.h - heightMap[x][y] : 0;
                    water += gap;
                    queue.offer(new Grid(x, y, heightMap[x][y] + gap));
                    visited[x][y] = true;
                }
            }

        }
        return water;
    }
}