星星博客 »  > 

【数据结构与算法】Leetcode407 接雨水 II -困难(堆、BFS)

来源:力扣

题目描述

给你一个 m x n 的矩阵,其中的值均为非负整数,代表二维高度图每个单元的高度,请计算图中形状最多能接多少体积的雨水。

示例:
给出如下 3x6 的高度图:
[
  [1,4,3,1,3,2],
  [3,2,1,3,2,4],
  [2,3,3,2,3,1]
]
返回 4

如上图所示,这是下雨前的高度图[[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]] 的状态。

下雨后,雨水将会被存储在这些方块中。总的接雨水量是4

提示:
1 <= m, n <= 110
0 <= heightMap[i][j] <= 20000

优先队列(堆)

  在看这题的解法前可以去我之前讲解接雨水I的那篇博客看一下,那里也应用了这里的解法。
  在那篇博客里我将双指针的本质做了讲解,即可以看作是从数组两端同时注水,然后从外到内,从低到高,直到左右两边的指针相遇即左右两边的水面相遇。
  现在用矩阵取代之前的数组,从一维扩展到了二维,之前的数组两端映射为矩阵的最外围的四边,接雨水I从数组两端注水,那么在这里则是从矩阵最外围的四边同时注水,从外到内,从低到高直至淹没所有位置。
  在接雨水I中可以用双指针指向数组两端,这里则没法为所有外围位置都分配一个指针,所有必须要用到接雨水I从双指针解法延伸出来的优先队列解法。和接雨水I中先将数组两端元素压入队列中类似,先把矩阵最外围四边的所有元素压入队列,然后依次弹出堆顶元素即新的水面值,依次访问上下左右与其相邻的位置,遇到比当前水面值低的就计算积水量,并将新遍历到的元素和新的水面值都压入队列中,重复这些操作直至队列为空,即所有位置都已经遍历过了,用visited数组记录所有遍历过的位置,防止重复计算。这里也用到了BFS广度优先搜索的思想。
  整个过程就像是随着水面不断升高,水不断地从四周向中间漫延一样,最终将所有柱子淹没,完成遍历,此时将水源撤走,留下的水即是矩阵最终的积水量。

代码(Python)

class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        # 如果矩阵的长或者宽小于等于2则不可能积水
        if len(heightMap) <= 2 or len(heightMap[0]) <= 2:
            return 0

        m, n = len(heightMap), len(heightMap[0])
        # visited用于储存已经遍历过的位置
        visited = set()
        # 优先队列/堆,用于从外到内、从小到大动态地遍历各个位置
        heap = []
        # 初始化堆,将矩阵外围的所有位置及其高度保存进堆中
        for i in range(m):
            heapq.heappush(heap, (heightMap[i][0], i, 0))
            heapq.heappush(heap, (heightMap[i][n-1], i, n-1))
            visited.add((i, 0))
            visited.add((i, n-1))
        for j in range(1, n-1):
            heapq.heappush(heap, (heightMap[0][j], 0, j))
            heapq.heappush(heap, (heightMap[m-1][j], m-1, j))
            visited.add((0, j))
            visited.add((m-1, j))
        
        ans = 0
        while heap:
            h, i, j = heapq.heappop(heap)
            # 遍历当前位置上下左右相邻的位置
            for x, y in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
                new_i, new_j = i + x, j + y
                # 判断新位置是否有效且没被访问
                if 0 <= new_i < m and 0 <= new_j < n and (new_i, new_j) not in visited:
                    near_h = heightMap[new_i][new_j]
                    # 遇到比水面低的,计算积水量
                    if near_h < h:
                        ans += h - near_h
                    # 将新位置及其高度与当前水面较高的那个存进堆中
                    heapq.heappush(heap, (max(h, near_h), new_i, new_j))
                    # 标记当前位置为已访问
                    visited.add((new_i, new_j))
        return ans

相关文章