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

from heapq import heappop, heappush
from typing import List

class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        if not heightMap or not heightMap[0]:
            return 0
        
        row, col = len(heightMap), len(heightMap[0])
        
        # Min-heap to keep track of the lowest boundary cells
        heap = []
        visited = [[False] * col for _ in range(row)]
        
        # Push all the boundary cells into the heap
        for r in range(row):
            for c in [0, col - 1]:  # Leftmost and rightmost columns
                heappush(heap, (heightMap[r][c], r, c))
                visited[r][c] = True
        for c in range(col):
            for r in [0, row - 1]:  # Topmost and bottommost rows
                heappush(heap, (heightMap[r][c], r, c))
                visited[r][c] = True

        directions = [(0, 1), (0, -1), (1, 0), (-1, 0)]  # Right, left, down, up
        total_water = 0

        # Process cells in the heap (starting from the lowest boundary cells)
        while heap:
            height, r, c = heappop(heap)
            
            # Visit all 4 directions
            for dr, dc in directions:
                nr, nc = r + dr, c + dc
                if 0 <= nr < row and 0 <= nc < col and not visited[nr][nc]:
                    # If this neighboring cell can trap water
                    total_water += max(0, height - heightMap[nr][nc])
                    
                    # Update the height of the neighboring cell to be the max of current height and its own
                    heappush(heap, (max(height, heightMap[nr][nc]), nr, nc))
                    visited[nr][nc] = True

        return total_water

19000