Lintcode 364. Trapping Rain Water II

Given n x m non-negative integers representing an elevation map 2d where the area of each cell is 1 x 1, compute how much water it is able to trap after raining.

Example

Given 5*4 matrix

[12,13,0,12]
[13,4,13,12]
[13,8,10,12]
[12,13,12,12]
[13,13,13,13]

return 14.

Code (Java)

public class Solution {
    /**
     * @param heights: a matrix of integers
     * @return: an integer
     */
    public int trapRainWater(int[][] heights) {
        if (heights == null || heights.length == 0) {
            return 0;
        }

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

        PriorityQueue<Pair> pq = new PriorityQueue<>(new MyPairComparator());
        boolean[][] visited = new boolean[m][n];

        // step 1: put borders into the pq
        //
        for (int i = 0; i < n; i++) {
            pq.offer(new Pair(0, i, heights[0][i]));
            visited[0][i] = true;

            pq.offer(new Pair(m - 1, i, heights[m - 1][i]));
            visited[m - 1][i] = true;
        }

        for (int i = 1; i < m - 1; i++) {
            pq.offer(new Pair(i, 0, heights[i][0]));
            visited[i][0] = true;

            pq.offer(new Pair(i, n - 1, heights[i][n - 1]));
            visited[i][n - 1] = true;
        }
        
        int ans = 0;
        int[][] dirs = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}};
        while (!pq.isEmpty()) {
            Pair top = pq.poll();

            for (int i = 0; i < 4; i++) {
                int nx = top.x + dirs[i][0];
                int ny = top.y + dirs[i][1];

                if (nx >= 0 && nx < m && ny >= 0 && ny < n && !visited[nx][ny]) {
                    visited[nx][ny] = true;
                    ans += Math.max(0, top.val - heights[nx][ny]);
                    pq.offer(new Pair(nx, ny, Math.max(top.val, heights[nx][ny])));
                }
            }
        }

        return ans;
    }
}

class Pair {
    int x;
    int y;
    int val;

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

class MyPairComparator implements Comparator<Pair> {
    @Override
    public int compare(Pair a, Pair b) {
        return a.val - b.val;
    }
}