Leetcode: 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.

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.

 1 public class Solution {
 2     public class Cell {
 3         int row;
 4         int col;
 5         int height;
 6         public Cell(int x, int y, int val) {
 7             this.row = x;
 8             this.col = y;
 9             this.height = val;
10         }
11     }
12     
13     public int trapRainWater(int[][] heightMap) {
14         if (heightMap==null || heightMap.length<=2 || heightMap[0].length<=2) return 0;
15         int m = heightMap.length;
16         int n = heightMap[0].length;
17         int res = 0;
18         PriorityQueue<Cell> queue = new PriorityQueue<Cell>(1, new Comparator<Cell>() {
19             public int compare(Cell c1, Cell c2) {
20                 return c1.height - c2.height;
21             }
22         });
23         HashSet<Integer> visited = new HashSet<Integer>();
24         for (int i=0; i<m; i++) {
25             queue.offer(new Cell(i, 0, heightMap[i][0]));
26             queue.offer(new Cell(i, n-1, heightMap[i][n-1]));
27             visited.add(i*n+0);
28             visited.add(i*n+n-1);
29         }
30         for (int j=0; j<n; j++) {
31             queue.offer(new Cell(0, j, heightMap[0][j]));
32             queue.offer(new Cell(m-1, j, heightMap[m-1][j]));
33             visited.add(0*n+j);
34             visited.add((m-1)*n+j);
35         }
36         int[][] directions = new int[][]{{-1, 0}, {1, 0}, {0, 1}, {0, -1}};
37         while (!queue.isEmpty()) {
38             Cell cur = queue.poll();
39             for (int[] dir : directions) {
40                 int row = cur.row + dir[0];
41                 int col = cur.col + dir[1];
42                 if (row>=0 && row<m && col>=0 && col<n && !visited.contains(row*n+col)) {
43                     visited.add(row*n+col);
44                     res += Math.max(0, cur.height - heightMap[row][col]);
45                     queue.offer(new Cell(row, col, Math.max(cur.height, heightMap[row][col])));
46                 }
47             }
48         }
49         return res;
50     }
51 }