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:

        m, n = len(heightMap), len(heightMap[0])

        hp = []

        visited = [[False for _ in range(n)] for _ in range(m)]

        # 最外围围栏入堆

        for i in range(m):

            for j in range(n):

                if i == 0 or j == 0 or i == m - 1 or j == n - 1:

                    # 最小堆，保证最矮的围栏出堆

                    heapq.heappush(hp, (heightMap[i][j], i, j))

                    visited[i][j] = True

        

        ans = 0

        while hp:

            h, r, c = heapq.heappop(hp)

            for nr, nc in ((r + 1, c), (r - 1, c), (r, c + 1), (r, c - 1)):

                if 0 <= nr < m and 0 <= nc < n and not visited[nr][nc]:

                    # 围栏比内部高，可以灌水

                    if h > heightMap[nr][nc]:

                        ans += h - heightMap[nr][nc]

                    # 忽略当前围栏，在(nr, nc)处新建围栏

                    visited[nr][nc] = True

                    # 新的围栏入堆

                    heapq.heappush(hp, (max(h, heightMap[nr][nc]), nr, nc))

        return ans