407. Trapping Rain Water II

 avatar
Atlas
c_cpp
2 months ago
1.7 kB
6
Indexable
#include <vector>
#include <queue>
#include <algorithm>
using namespace std;

class Solution {
public:
    int trapRainWater(vector<vector<int>>& heightMap) {
        int m = heightMap.size();
        int n = heightMap[0].size();
        if (m < 3 || n < 3) return 0; // 最小要 3x3 才能蓄水

        using T = tuple<int, int, int>; // (高度, x, y)
        priority_queue<T, vector<T>, greater<T>> pq;
        vector<vector<bool>> visited(m, vector<bool>(n, false));

        // Step 1: 將邊界推入 PQ,並標記 visited
        for (int i = 0; i < m; ++i) {
            pq.emplace(heightMap[i][0], i, 0);
            pq.emplace(heightMap[i][n - 1], i, n - 1);
            visited[i][0] = visited[i][n - 1] = true;
        }
        for (int j = 1; j < n - 1; ++j) {
            pq.emplace(heightMap[0][j], 0, j);
            pq.emplace(heightMap[m - 1][j], m - 1, j);
            visited[0][j] = visited[m - 1][j] = true;
        }

        int totalWater = 0;
        vector<int> dir = {0, 1, 0, -1, 0};

        // Step 2: BFS 擴展
        while (!pq.empty()) {
            auto [h, x, y] = pq.top(); pq.pop();
            for (int d = 0; d < 4; ++d) {
                int nx = x + dir[d], ny = y + dir[d + 1];
                if (nx < 0 || ny < 0 || nx >= m || ny >= n || visited[nx][ny]) continue;
                visited[nx][ny] = true;
                // 如果 neighbor 比堤壩矮,就可以接水
                totalWater += max(0, h - heightMap[nx][ny]);
                // 將 neighbor 當作新的堤壩(高度要取 max)
                pq.emplace(max(h, heightMap[nx][ny]), nx, ny);
            }
        }

        return totalWater;
    }
};
Editor is loading...
Leave a Comment