[LeetCode] 407. Trapping Rain Water II

收集雨水之二。题意跟42题很接近，但是是从二维的变成三维了。思路参见grandyang大神的帖子。这个题我用Java实现的。

时间O(m * n * log(m + n))

空间O(m * n)

 1 class Solution {
 2     int[][] dirs = new int[][] { { 1, 0 }, { -1, 0 }, { 0, 1 }, { 0, -1 } };
 3 
 4     public int trapRainWater(int[][] heightMap) {
 5         int m = heightMap.length;
 6         int n = (m == 0 ? 0 : heightMap[0].length);
 7         int res = 0;
 8 
 9         PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> a[2] - b[2]);
10         boolean[][] visited = new boolean[m][n];
11         for (int i = 0; i < m; i++) {
12             pq.offer(new int[] { i, 0, heightMap[i][0] });
13             pq.offer(new int[] { i, n - 1, heightMap[i][n - 1] });
14             visited[i][0] = visited[i][n - 1] = true;
15         }
16         for (int j = 1; j < n - 1; j++) {
17             pq.offer(new int[] { 0, j, heightMap[0][j] });
18             pq.offer(new int[] { m - 1, j, heightMap[m - 1][j] });
19             visited[0][j] = visited[m - 1][j] = true;
20         }
21 
22         while (!pq.isEmpty()) {
23             int[] cell = pq.poll();
24             for (int[] d : dirs) {
25                 int i = cell[0] + d[0], j = cell[1] + d[1];
26                 if (i < 0 || i >= m || j < 0 || j >= n || visited[i][j]) {
27                     continue;
28                 }
29                 res += Math.max(0, cell[2] - heightMap[i][j]);
30                 pq.offer(new int[] { i, j, Math.max(heightMap[i][j], cell[2]) });
31                 visited[i][j] = true;
32             }
33         }
34         return res;
35     }
36 }