<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:
		# number of rows
		m = len(heightMap)
		# number of columns
		n = len(heightMap[0])
		#initialize visited nodes and heapq
		flag = [[1 for _ in range(n)] for _ in range(m)]
		minheap = []
		for i in range(m):
			for j in range(n):
				#first add to the minheap the border of the enclosed rain water area
				if i == 0 or j == 0 or i == m - 1 or j == n - 1:
					# push into the heap (height, i,j) since we need to know the index of each height for later
					heapq.heappush(minheap,(heightMap[i][j],i,j))
				else:
					flag[i][j] = 0
		#directional vector
		dx = [1,0,-1,0]
		dy = [0,1,0,-1]
		ans = 0
		while minheap:
			#the height of the water is determined by the lowest point in the border
			curr = heapq.heappop(minheap)
			for i in range(4):
				newi = curr[1] + dx[i]
				newj = curr[2] + dy[i]
				# check out of bound condition or if the node is already visited
				if newi < 0 or newi >= m or newj < 0 or newj >= n or flag[newi][newj]:
					continue
				# if the new location has a smaller height than the current one (being popped), then it can at least accumulate as much rain water as the current location
				if heightMap[newi][newj] <= curr[0]:
					#update the new border with new height (after accumulating water)
					heapq.heappush(minheap, (curr[0], newi, newj))
				else:
					heapq.heappush(minheap, (heightMap[newi][newj], newi, newj))
				#mark the visited node
				flag[newi][newj] = 1
				ans += max(0, curr[0] - heightMap[newi][newj])
		return ans
	
18900