16. Trapping Rain Water II

Link to the problem

Description

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.

Idea

Maintains the boundary of the map. Every time, remove the shortest block, and add all its new neighbors. Any neighbor shorter than the removed block can hold the water between itself and the removed block.

Solution

struct Comparator {
    const bool operator() (const vector<int> &v1, const vector<int> &v2) {
        return v1[2] > v2[2];
    }
};

class Solution {
public:
    int trapRainWater(vector<vector<int> >& heightMap) {
        priority_queue<vector<int>, vector<vector<int> >, Comparator> pq;
        if (heightMap.empty()) return 0;
        int m = heightMap.size(), n = heightMap[0].size();
        if (m == 1 || n <= 1) return 0;
        unordered_set<string> visited;
        // Enqueue all boundary
        for (int i = 0; i < m; i++) {
            pq.push(vector<int> {i, 0, heightMap[i][0]});
            pq.push(vector<int> {i, n - 1, heightMap[i][n - 1]});
            visited.insert(to_string(i) + "," + to_string(0));
            visited.insert(to_string(i) + "," + to_string(n - 1));
        }
        for (int i = 1; i < n - 1; i++) {
            pq.push(vector<int> {0, i, heightMap[0][i]});
            pq.push(vector<int> {m - 1, i, heightMap[m - 1][i]});
            visited.insert(to_string(0) + "," + to_string(i));
            visited.insert(to_string(m - 1) + "," + to_string(i));
        }
        int trapped = 0;
        vector<vector<int> > diff = {{0, 1}, {1, 0}, {0, -1}, {-1, 0}};
        while (!pq.empty()) {
            int curr_row = pq.top()[0], curr_col = pq.top()[1], curr_height = pq.top()[2];
            pq.pop();
            for (auto it = diff.begin(); it != diff.end(); it++) {
                int row = curr_row + it->at(0), col = curr_col + it->at(1);
                string key = to_string(row) + "," + to_string(col);
                if (0 <= row && row < m && 0 <= col && col < n &&
                   visited.find(key) == visited.end()) {
                    if (curr_height > heightMap[row][col]) {
                        trapped += curr_height - heightMap[row][col];
                    }
                    pq.push(vector<int> {row, col, max(heightMap[row][col], curr_height)});
                    visited.insert(key);
                }
            }
        }
        return trapped;
    }
};

40 / 40 test cases passed.
Runtime: 109 ms