leetcode407接雨水II

class Solution {
    #define P pair<int, int>
public:
    // 最小堆
    int dir[5] = {-1, 0, 1, 0, -1};
    int trapRainWater(vector<vector<int>>& heightMap) {
        priority_queue<P, vector<P>, greater<P> > q;
        int n = heightMap.size(), m = heightMap[0].size(), ans = 0;
        vector<vector<int>> vis(n, vector<int>(m));
        for(int i = 0; i < n; i++){
            for(int j = 0; j < m; j++){
                if(i == 0 || i == n  - 1 || j == 0 || j == m - 1){
                    q.push(make_pair(heightMap[i][j], i * m + j));
                    vis[i][j] = 1;
                }
            }
        }
        while(!q.empty()){
            P p = q.top(); q.pop();
            for(int i = 0; i < 4; i++){
                int x = p.second / m + dir[i];
                int y = p.second % m + dir[i + 1];
                //cout <<x << " " << y <<endl;
                if(x < 0 || x >= n || y < 0 || y >= m || vis[x][y]) continue;
                ans += max(0, p.first - heightMap[x][y]);
                //cout << ans << " " << endl;
                vis[x][y] = 1;
                q.push(make_pair(max(heightMap[x][y], p.first), x * m + y));
            }

        }
        return ans;
    }
};