[LintCode] Trapping Rain Water II

Trapping Rain Water II

Given n x m non-negative integers representing an elevation map 2d where the area of each cell is 1 x 1, compute how much water it is able to trap after raining.

 

Example

Given 5*4 matrix

[12,13,0,12]
[13,4,13,12]
[13,8,10,12]
[12,13,12,12]
[13,13,13,13]

return 14.

Tags Expand 

LintCode Copyright Heap Matrix

最小化堆。注意在往堆里放新元素的时候，要同时更新位置的水位高度。即：

heap.emplace(xx, yy, max(u.h, heights[xx][yy]));

还要注意优先队列默认是大顶堆，默认使用<号，所以要想使用小顶堆，在自己写cmp仿函数时，要返回a > b, 还不是a < b。

class Solution {
public:
    /**
     * @param heights: a matrix of integers
     * @return: an integer
     */
    struct cell {
        int x, y;
        int h;
        cell() {}
        cell(int _x, int _y, int _h) : x(_x), y(_y), h(_h) {}
    };
    struct cmp {
        bool operator() (const cell &a, const cell &b) {
            return a.h > b.h;
        }
    };
    int trapRainWater(vector<vector<int> > &heights) {
        // write your code here
        if (heights.empty() || heights[0].empty()) return 0;
        int n = heights.size(), m = heights[0].size();
        vector<vector<bool>> visit(n, vector<bool>(m, false));
        priority_queue<cell, vector<cell>, cmp> heap;
        for (int i = 0; i < n; ++i) {
            heap.emplace(i, 0, heights[i][0]);
            heap.emplace(i, m - 1, heights[i][m-1]);
            visit[i][0] = visit[i][m-1] = true;
        }
        for (int j = 0; j < m; ++j) {
            heap.emplace(0, j, heights[0][j]);
            heap.emplace(n - 1, j, heights[n-1][j]);
            visit[0][j] = visit[n-1][j] = true;
        }
        int res = 0;
        const int dx[4] = {0, 1, 0, -1};
        const int dy[4] = {1, 0, -1, 0};
        while (!heap.empty()) {
            auto u = heap.top();
            heap.pop();
            for (int i = 0; i < 4; ++i) {
                int xx = u.x + dx[i], yy = u.y + dy[i];
                if (xx >= 0 && xx < n && yy >= 0 && yy < m && !visit[xx][yy]) {
                    res += max(0, u.h - heights[xx][yy]);
                    heap.emplace(xx, yy, max(u.h, heights[xx][yy]));
                    visit[xx][yy] = true;
                }
            }
        }
        return res;
    }
};