[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。

 1 class Solution {
 2 public:
 3     /**
 4      * @param heights: a matrix of integers
 5      * @return: an integer
 6      */
 7     struct cell {
 8         int x, y;
 9         int h;
10         cell() {}
11         cell(int _x, int _y, int _h) : x(_x), y(_y), h(_h) {}
12     };
13     struct cmp {
14         bool operator() (const cell &a, const cell &b) {
15             return a.h > b.h;
16         }
17     };
18     int trapRainWater(vector<vector<int> > &heights) {
19         // write your code here
20         if (heights.empty() || heights[0].empty()) return 0;
21         int n = heights.size(), m = heights[0].size();
22         vector<vector<bool>> visit(n, vector<bool>(m, false));
23         priority_queue<cell, vector<cell>, cmp> heap;
24         for (int i = 0; i < n; ++i) {
25             heap.emplace(i, 0, heights[i][0]);
26             heap.emplace(i, m - 1, heights[i][m-1]);
27             visit[i][0] = visit[i][m-1] = true;
28         }
29         for (int j = 0; j < m; ++j) {
30             heap.emplace(0, j, heights[0][j]);
31             heap.emplace(n - 1, j, heights[n-1][j]);
32             visit[0][j] = visit[n-1][j] = true;
33         }
34         int res = 0;
35         const int dx[4] = {0, 1, 0, -1};
36         const int dy[4] = {1, 0, -1, 0};
37         while (!heap.empty()) {
38             auto u = heap.top();
39             heap.pop();
40             for (int i = 0; i < 4; ++i) {
41                 int xx = u.x + dx[i], yy = u.y + dy[i];
42                 if (xx >= 0 && xx < n && yy >= 0 && yy < m && !visit[xx][yy]) {
43                     res += max(0, u.h - heights[xx][yy]);
44                     heap.emplace(xx, yy, max(u.h, heights[xx][yy]));
45                     visit[xx][yy] = true;
46                 }
47             }
48         }
49         return res;
50     }
51 };