<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>

class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        # the rain cannot be trapped at the edge matrix
        # so we will add them to the heap to go so we can iterate through

        # then at each of the coordinate
        # if we found a value higher than the current 'wall'
        # we add it to the heap
        
        # and if there is a value that's lower than the current 'wall'
        # it can contain water

        if not heightMap or not heightMap[0]:
            return 0
        
        rows = len(heightMap)
        cols = len(heightMap[0])

        # we can't trap water if there is only 2 rows/cols
        if rows < 3 or cols < 3:
            return 0
        

        heap = []
        for r in range(rows):
            for c in range(cols):
                if r == 0 or r == rows - 1 or c == 0 or c == cols - 1:
                    heapq.heappush(heap, (heightMap[r][c], r, c))
                    # mark as visited
                    heightMap[r][c] = -1
        
        level = 0
        result = 0

        while heap:
            curHeight, r, c = heapq.heappop(heap)
            level = max(curHeight, level)

            for newR, newC in [(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)]:
                if 0 <= newR < rows and 0 <= newC < cols and heightMap[newR][newC] != -1:
                    heapq.heappush(heap, (heightMap[newR][newC], newR, newC))

                    if heightMap[newR][newC] < level:
                        result += level - heightMap[newR][newC]
                    
                    heightMap[newR][newC] = -1
        
        return result

18500