407. Trapping Rain Water II
Given an m x n matrix of positive integers representing the height of each unit cell in a 2D elevation map, compute the volume of water it is able to trap after raining.
Note:
Both m and n are less than 110. The height of each unit cell is greater than 0 and is less than 20,000.
Example:
Given the following 3x6 height map:
[
[1,4,3,1,3,2],
[3,2,1,3,2,4],
[2,3,3,2,3,1]
]
Return 4.
The above image represents the elevation map [[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]] before the rain.
After the rain, water are trapped between the blocks. The total volume of water trapped is 4.
Solution: Heap
import heapq
class Solution(object):
def trapRainWater(self, heights):
"""
:type heights: List[List[int]]
:rtype: int
"""
if not heights:
return 0
rain = 0
rows = len(heights)
cols = len(heights[0])
visited = [[False] * cols for _ in xrange(rows)]
q = []
for row, col in self.border(rows, cols):
visited[row][col] = True
heapq.heappush(q, (heights[row][col], row, col))
while q:
height, row, col = heapq.heappop(q)
for i, j in self.neighbors(row, col, rows, cols):
if not visited[i][j]:
visited[i][j] = True
if heights[i][j] < height:
rain += height - heights[i][j]
heights[i][j] = height
heapq.heappush(q, (heights[i][j], i, j))
return rain
def neighbors(self, row, col, rows, cols):
for x, y in (1, 0), (-1, 0), (0, 1), (0, -1):
i, j = row + x, col + y
if 0 <= i < rows and 0 <= j < cols:
yield i, j
def border(self, rows, cols):
for col in xrange(0, cols):
yield 0, col
for row in xrange(1, rows):
yield row, cols - 1
if rows > 1 and cols > 1:
for col in xrange(cols - 2, -1, -1):
yield rows - 1, col
for row in xrange(rows - 2, 0, -1):
yield row, 0
Lessons:
- Once fill water in a cell, set height of the cell to the water level, and push this height to heap.