<p>Given an <code>m x n</code> integer matrix <code>heightMap</code> representing the height of each unit cell in a 2D elevation map, return <em>the volume of water it can trap after raining</em>.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2021/04/08/trap1-3d.jpg" style="width: 361px; height: 321px;" />
<pre>
<strong>Input:</strong> heightMap = [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]]
<strong>Output:</strong> 4
<strong>Explanation:</strong> After the rain, water is trapped between the blocks.
We have two small ponds 1 and 3 units trapped.
The total volume of water trapped is 4.
</pre>

<p><strong class="example">Example 2:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2021/04/08/trap2-3d.jpg" style="width: 401px; height: 321px;" />
<pre>
<strong>Input:</strong> 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]]
<strong>Output:</strong> 10
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>m == heightMap.length</code></li>
	<li><code>n == heightMap[i].length</code></li>
	<li><code>1 &lt;= m, n &lt;= 200</code></li>
	<li><code>0 &lt;= heightMap[i][j] &lt;= 2 * 10<sup>4</sup></code></li>
</ul>

import heapq

class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        # w[x,y] = max(h[x,y], min(w[x+1,y], w[x-1,y], w[x,y+1], w[x,y-1]))
        heap = []
        m, n = len(heightMap), len(heightMap[0])
        levelMap = [[None for _ in range(n)] for _ in range(m)]
        for i in [0, m - 1]:
            for j in range(n):
                level = heightMap[i][j]
                heapq.heappush(heap, (level, i, j))
        for i in range(1, m - 1):
            for j in [0, n - 1]:
                level = heightMap[i][j]
                heapq.heappush(heap, (level, i, j))
        while heap:
            level, i0, j0 = heapq.heappop(heap)
            # print("heap", level, i0, j0)
            stack = [(i0, j0)]
            while stack:
                i0, j0 = stack.pop()
                # print("stack1", level, i0, j0)
                if levelMap[i0][j0] is not None:
                    continue
                # print("stack2", level, i0, j0)
                h = heightMap[i0][j0]
                # print("h", h, i0, j0)
                if h > level:
                    heapq.heappush(heap, (h, i0, j0))
                    continue
                levelMap[i0][j0] = level
                for i, j in [(i0, j0 - 1), (i0, j0 + 1), (i0 - 1, j0), (i0 + 1, j0)]:
                    if (0 <= i < m) and (0 <= j < n) and levelMap[i][j] is None:
                        # print("append", i, j)
                        stack.append((i, j))
        return sum(sum(levelMap[i][j] - heightMap[i][j] for j in range(n)) for i in range(m))
21500