407. 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 is trapped between the blocks. The total volume of water trapped is 4.

// BFS + priority_queue
int trapRainWater(vector<vector<int>>& heightMap) { // time: O(mn*log(mn)); space: O(mn)
    if (heightMap.empty()) return 0;
    int res = 0, seaLevel = 0, m = heightMap.size(), n = heightMap[0].size();
    priority_queue<pair<int, int>, vector<pair<int, int> >, greater<pair<int, int> > > pq;
    vector<vector<int> > dirs({{-1, 0}, {1, 0}, {0, -1}, {0, 1}});
    vector<vector<bool> > visited(m, vector<bool> (n, false));
    // Add boundaries
    for (int i = 0; i < m; ++i) {
        for (int j = 0; j < n; ++j) {
            if (i == 0 || i == m - 1 || j == 0 || j == n - 1) {
                pq.push({heightMap[i][j], i * n + j}); // coordinates compression
                visited[i][j] = true;
            }
        }
    }
    while (!pq.empty()) {
        auto t = pq.top(); pq.pop();
        int h = t.first, r = t.second / n, c = t.second % n;
        seaLevel = max(seaLevel, h);
        for (auto& dir : dirs) {
            int x = r + dir[0], y = c + dir[1];
            if (x < 0 || x >= m || y < 0 || y >= n || visited[x][y]) continue;
            visited[x][y] = true;
            if (heightMap[x][y] < seaLevel) res += (seaLevel - heightMap[x][y]);
            pq.push({heightMap[x][y], x * n + y});
        }
    }
    return res;
}