给你一个 m x n 的矩阵,其中的值均为非负整数,代表二维高度图每个单元的高度,请计算图中形状最多能接多少体积的雨水。
输入: heightMap = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]
输出: 4
解释: 下雨后,雨水将会被上图蓝色的方块中。总的接雨水量为1+2+1=4。
class Solution:
def trapRainWater(self, heightMap: List[List[int]]) -> int:
import heapq
if not heightMap: return 0
heap = []
cur_max = float("-inf")
visited = set()
row = len(heightMap)
col = len(heightMap[0])
res = 0
# 边界放进去
# 行
for j in range(col):
heapq.heappush(heap, [heightMap[0][j], 0, j])
heapq.heappush(heap, [heightMap[row - 1][j], row - 1, j])
visited.add((0, j))
visited.add((row - 1, j))
# 列
for i in range(row):
heapq.heappush(heap, [heightMap[i][0], i, 0])
heapq.heappush(heap, [heightMap[i][col - 1], i, col - 1])
visited.add((i, 0))
visited.add((i, col - 1))
while heap:
h, i, j = heapq.heappop(heap)
cur_max = max(h, cur_max)
for x, y in [[-1, 0], [1, 0], [0, 1], [0, -1]]:
tmp_i = i + x
tmp_j = j + y
if tmp_i < 0 or tmp_i >= row or tmp_j < 0 or tmp_j >= col or (tmp_i, tmp_j) in visited:
continue
visited.add((tmp_i, tmp_j))
if heightMap[tmp_i][tmp_j] < cur_max:
res += cur_max - heightMap[tmp_i][tmp_j]
heapq.heappush(heap, [heightMap[tmp_i][tmp_j], tmp_i, tmp_j])
return res