Trapping Rain Water II.cpp

class Solution {
public:
    int trapRainWater(vector<vector<int>>& heightMap) {
        int m = heightMap.size(), n = m? heightMap[0].size(): 0;
        auto comp = [&](const int i, const int j){ return heightMap[i / n][i % n] > heightMap[j / n][j % n]; };
        priority_queue<int, vector<int>, decltype(comp)> pq(comp);
        vector<pair<int, int>> dirs = {{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
        vector<bool> visited(m * n, false);
        for(int i = 0; i < m; ++i)
        {
            pq.push(i * n);
            pq.push(i * n + n - 1);
            visited[i * n] = visited[i * n + n - 1] = true;
        }
        for(int i = 0; i < n; ++i)
        {
            pq.push(i);
            pq.push((m - 1) * n + i);
            visited[i] = visited[(m - 1) * n + i] = true;
        }
        int vol = 0;

        while(pq.size())
        {
            int curr = pq.top(), i = curr / n, j = curr % n;
            pq.pop();
            for(auto& dir : dirs)
            {
                int x = i + dir.first, y = j + dir.second;
                if(x >= 0 && x < m && y >= 0 && y < n && !visited[x * n + y])
                {
                    visited[x * n + y] = true;
                    vol += heightMap[i][j] > heightMap[x][y]? heightMap[i][j] - heightMap[x][y]: 0;
                    heightMap[x][y] = max(heightMap[i][j], heightMap[x][y]);
                    pq.push(x * n + y);
                }
            }
        }
        return vol;
    }
};