LeetCode 407: Trapping Rain Water II

Note:

The update height with MAX(current height, real height) is because how many water this block can held depends on all outside heights.

Another thought is that when we traversed to this block, we found that it has been filled with water, then we make it height as the max(min(all neighbors), current height) to make it flat.

class Solution {
    class Block {
        int x;
        int y;
        int height;
        public Block(int x, int y, int height) {
            this.x = x;
            this.y = y;
            this.height = height;
        }
    }
    public int trapRainWater(int[][] heightMap) {
        if (heightMap.length < 3) return 0;
        PriorityQueue<Block> queue = new PriorityQueue<>((b1, b2) -> b1.height - b2.height);
        boolean[][] visited = new boolean[heightMap.length][heightMap[0].length];
        for (int i = 0; i < heightMap.length; i++) {
            visited[i][0] = true;
            visited[i][heightMap[0].length - 1] = true;
            queue.offer(new Block(i, 0, heightMap[i][0]));
            queue.offer(new Block(i, heightMap[0].length - 1, heightMap[i][heightMap[0].length - 1]));
        }
        
        for (int i = 0; i < heightMap[0].length; i++) {
            visited[0][i] = true;
            visited[heightMap.length - 1][i] = true;
            queue.offer(new Block(0, i, heightMap[0][i]));
            queue.offer(new Block(heightMap.length - 1, i, heightMap[heightMap.length - 1][i]));
        }
        
        int[][] dirs = new int[][] {{0, 1}, {0, -1}, {1, 0}, {-1, 0}};
        
        int result = 0;
        while (!queue.isEmpty()) {
            Block current = queue.poll();
            for (int[] dir : dirs) {
                int x = current.x + dir[0];
                int y = current.y + dir[1];
                if (x >= 0 && x < heightMap.length && y >= 0 && y < heightMap[0].length && !visited[x][y]) {
                    visited[x][y] = true;
                    result += Math.max(0, current.height - heightMap[x][y]);
                    queue.offer(new Block(x, y, Math.max(heightMap[x][y], current.height)));
                }
            }
        }
        return result;
    }
}