Leetcode - Trapping Rain Water II

My code:

public class Solution {
    private class Cell {
        int x;
        int y;
        int h;
        Cell(int x, int y, int h) {
            this.x = x;
            this.y = y;
            this.h = h;
        }
    }
    
    public int trapRainWater(int[][] heightMap) {
        if (heightMap == null || heightMap.length <= 2 || heightMap[0].length <= 2) {
            return 0;
        }
        
        int row = heightMap.length;
        int col = heightMap[0].length;
        boolean[][] mark = new boolean[row][col];
        
        PriorityQueue pq = new PriorityQueue(1, new Comparator() {
            public int compare(Cell c1, Cell c2) {
                return c1.h - c2.h;
            }    
        });
        
        for (int i = 0; i < col; i++) {
            mark[0][i] = true;
            mark[row - 1][i] = true;
            pq.offer(new Cell(0, i, heightMap[0][i]));
            pq.offer(new Cell(row - 1, i, heightMap[row - 1][i]));
        }        
        
        for (int i = 0; i < row; i++) {
            mark[i][0] = true;
            mark[i][col - 1] = true;
            pq.offer(new Cell(i, 0, heightMap[i][0]));
            pq.offer(new Cell(i, col - 1, heightMap[i][col - 1]));
        }
        
        int ret = 0;
        int[][] dir = new int[][]{{0, -1}, {0, 1}, {-1, 0}, {1, 0}};
        while (!pq.isEmpty()) {
            Cell c = pq.poll();
            for (int i = 0; i < 4; i++) {
                int nei_x = c.x + dir[i][0];
                int nei_y = c.y + dir[i][1];
                if (nei_x >= 0 && nei_x < row && nei_y >= 0 && nei_y < col && !mark[nei_x][nei_y]) {
                    mark[nei_x][nei_y] = true;
                    ret += Math.max(0, c.h - heightMap[nei_x][nei_y]);
                    pq.offer(new Cell(nei_x, nei_y, Math.max(c.h, heightMap[nei_x][nei_y])));
                }
            }
        }
        
        return ret;
    }
}

reference:
https://discuss.leetcode.com/topic/60418/java-solution-using-priorityqueue/2

上面的链接讲得挺清楚的。

Anyway, Good luck, Richardo! -- 10/14/2016