政子的博客

技术|学习|随笔

LeetCode-Trapping Rain Water II

https://leetcode-cn.com/problems/trapping-rain-water-ii/

参考了 solution How to get the solution to 2-D “Trapping Rain Water” problem from 1-D case 中的方法，可以从上一题Trapping Rain Water 的思路中进行拓展。

首先，解题的关键是优先级队列，它能够以 O(nlogn)的时间复杂度来得到队列中的最小值。在 Trapping Rain Water 中，我们使用了两个指针，一个从左向右，一个从右向左，找到两个指针指向的较小高度来判断是否有积水，那么在这道题里，也是这样的思路，但不能使用两个指针，而只能使用一个优先级队列了。用优先级队列找到当前最小的高度，然后探查该高度四周会不会有积水，然后把这个高度四周的高度继续加入队列，继续观察。

大致思路明了以后，代码书写就会比较方便了：

  • 优先队列中包含三个元素：某位置的 xy 坐标以及其被探查到时的最大高度
  • 将二维数组四周的高度都加入优先级队列中作为初始条件（在第二次循环的时候注意不要将两边的值加入两次）
  • 加入优先级队列中的值都要在 visited 数组中设置为 true，防止之后重新被加入
  • 当队列不为空的时候，先得到最小的高度，然后探查其周围，并将周围的元素也加入队列中，注意要保持最高高度，如果周围高度有比队列中的最小高度小的，则将高度差加入 result
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
class Solution {
private:
    struct cmp{
        bool operator ()(vector<int>& a, vector<int>& b){
            return a[2]>b[2];
        }
    };
    vector<vector<int>> directions{
        vector<int>{-1,0},
        vector<int>{1,0},
        vector<int>{0,-1},
        vector<int>{0,1}
    };
public:
    int trapRainWater(vector<vector<int>>& heightMap) {
        int m=heightMap.size();
        if(m<=2) return 0;
        int n=heightMap[0].size();
        if(n<=2) return 0;
        int result=0;
        vector<vector<bool>> visited(m,vector<bool>(n,false));
        priority_queue<vector<int>,vector<vector<int>>,cmp> pq;
        for(int i=0;i<m;++i){
            pq.push(vector<int>{i,0,heightMap[i][0]});
            visited[i][0]=true;
            pq.push(vector<int>{i,n-1,heightMap[i][n-1]});
            visited[i][n-1]=true;
        }
        for(int j=1;j<n-1;++j){
            pq.push(vector<int>{0,j,heightMap[0][j]});
            visited[0][j]=true;
            pq.push(vector<int>{m-1,j,heightMap[m-1][j]});
            visited[m-1][j]=true;
        }
        while(!pq.empty()){
            auto cur=pq.top();
            pq.pop();
            for(auto d:directions){
                int i=cur[0]+d[0];
                int j=cur[1]+d[1];
                if(i<0 || i>m-1 || j<0 || j>n-1 || visited[i][j]) continue;
                if(heightMap[i][j]<cur[2]){
                    result+=cur[2]-heightMap[i][j];
                    pq.push(vector<int>{i,j,cur[2]});
                }else{
                    pq.push(vector<int>{i,j,heightMap[i][j]});
                }
                visited[i][j]=true;
            }
        }
        return result;
    }
};