LeetCode刷题Python之407. 接雨水 II

给你一个 m x n 的矩阵，其中的值均为非负整数，代表二维高度图每个单元的高度，请计算图中形状最多能接多少体积的雨水。

示例 1:

输入: heightMap = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]
输出: 4
解释: 下雨后，雨水将会被上图蓝色的方块中。总的接雨水量为1+2+1=4。
示例 2:

输入: heightMap = [[3,3,3,3,3],[3,2,2,2,3],[3,2,1,2,3],[3,2,2,2,3],[3,3,3,3,3]]
输出: 10
 

提示:

m == heightMap.length
n == heightMap[i].length
1 <= m, n <= 200
0 <= heightMap[i][j] <= 2 * 104

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

力扣https://leetcode-cn.com/problems/trapping-rain-water-ii/solution/jie-yu-shui-ii-by-leetcode-solution-vlj3/