1
import heapq
2

3

4
class Solution:
5
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
6
        ROW, COL = len(heightMap), len(heightMap[0])
7

8
        pq = []
9
        heapq.heapify(pq)
10
        visited = {}
11

12
        for row in range(ROW):
13
            for col in range(COL):
14
                if row == 0 or row == ROW - 1 or col == 0 or col == COL - 1:
15
                    heapq.heappush(pq, (heightMap[row][col], row, col))
16
                    visited[(row, col)] = True
17

18
        def getnbr(row, col):
19
            res = []
20
            if row - 1 >= 0:
21
                res.append((row - 1, col))
22
            if col - 1 >= 0:
23
                res.append((row, col - 1))
24
            if row + 1 < ROW:
25
                res.append((row + 1, col))
26
            if col + 1 < COL:
27
                res.append((row, col + 1))
28

29
            return res
30

31
        res = 0
32

33
        while pq:
34
            h, i, j = heapq.heappop(pq)
35

36
            for dx, dy in getnbr(i, j):
37
                if (dx, dy) not in visited:
38

39
                    res += max(0, h - heightMap[dx][dy])
40

41
                    heapq.heappush(pq, (max(h, heightMap[dx][dy]), dx, dy))
42
                    visited[(dx, dy)] = True
43

44
        return res