[LeetCode] 407. Trapping Rain Water II_hard tag: Heap

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.

这个是[LeetCode] 42. Trapping Rain Water_hard tag: Two Pointers的follow up，变成了二维数组，我们还是利用从外向内灌水的思想。

1. 遍历最外层的点，将最外层的点加入到visited set 和heap（python中是min heap）中

2. 找到heap中的最短板，从heap中pop出来，因为有上下左右四个方向，找到未走过的点neigh，因为外层已经标记过了，所以没走过的肯定是里面的点

3. ans += minHeap - neigh[height] if minHeap > neigh[height] else 0，再将neigh存到visited中，push进heap

4. 不停重复2，3，不停更新最小的点，直到heap为空

5. 返回ans

 

Code：

import heapq
class Solution:
    def trapWater2(self, heightMap):
        if not heightMap or not heightMap[0]: return 0
        ans, heap, visited, m, n = 0, [], set(), len(heightMap), len(heightMap[0])
        directions = [(0,1), (0, -1), (-1, 0), (1, 0)]
        for i in range(m):
            for j in range(n):
                if i == 0 or i == m - 1 or j == 0 or j == n - 1:
                    visited.add((i, j))
                    heapq.heappush(heap, (heightMap[i][j], i, j))
        while heap:
            height, r, c = heapq.heappop(heap)
            for i, j in directions:
                x = r + i
                y = c + j
                if 0<= x < m and 0 <= y < n and (x, y) not in visited:
                    ans += max(0, height - heightMap[x][y])
                    visited.add(x, y)
                    heapq.heappush(heap, (max(height, heightMap[x][y]), x, y))
        return ans