Trapping rain water 2

Question

Thoughts

这一题是1的follow up,将1d的问题变成了2d,通过分析,对于matrix的能够灌水的高度,取决于最外围一圈的高度中的最小值,而找到这个最小值我们可以用heap来快速得到,python的heap是用heapq这个library来实现

  1. 用heap保存最外围的点,并且快速求得最小值
  2. 用一个visited matrix保存visit过的点

每次得到最外围的最小值,然后往内灌水(四个方向)

Solution

from heapq import *
class Solution:
    # @param heights: a matrix of integers
    # @return: an integer
    def trapRainWater(self, heights):
        # write your code here
        if not heights:
            return 0

        heap = []
        rowNum = len(heights)
        colNum = len(heights[0])
        visited = [[False for a in range(colNum)] for b in range(rowNum)]
        result = 0

        #add outside into heap
        for i in range(rowNum):
            heappush(heap, (heights[i][0], i, 0))
            heappush(heap, (heights[i][colNum-1], i, colNum-1))
            visited[i][0] = True
            visited[i][colNum-1] = True

        for j in range(1, colNum):
            heappush(heap, (heights[0][j], 0, j))
            heappush(heap, (heights[rowNum-1][j], rowNum-1, j))
            visited[0][j] = True
            visited[rowNum-1][j] = True

        dx = [1, -1, 0, 0]
        dy = [0, 0, -1, 1]
        while heap:
            curr = heappop(heap)
            for x in range(4):
                nx = curr[1] + dx[x]
                ny = curr[2] + dy[x]

                if nx >= 0 and nx < rowNum and ny >= 0 and ny < colNum and not visited[nx][ny]:
                    visited[nx][ny] = True
                    heappush(heap, (max(curr[0], heights[nx][ny]), nx, ny))
                    result += max(0, curr[0] - heights[nx][ny])

        return result