【算法】经典算法题：接雨水 I & II（Python实现）

目录

一、接雨水 I

1、动态规划

（1）算法原理

（2）Python实现

2、双指针

（1）算法原理

（2）Python实现

 二、接雨水 II

1、最小堆：二维情况（同接雨水 I）

（1）算法原理

（2）Python实现

2、最小堆：三维情况

（1）算法原理

（2）Python实现

---

一、接雨水 I

leetcode.42 链接：https://leetcode.cn/problems/trapping-rain-water/

难度：hard

这里介绍了两种方法：动态规划（时间和空间复杂度都为O(n)），双指针（时间复杂度O(n)，空间复杂度O(1)）。

1、动态规划

时间复杂度：O(n)

空间复杂度：O(n)

（1）算法原理

• 首先，维护两个数组：
          leftMax：第i个元素表示第i个柱子及左边位置中，height的最大高度。
          rightMax：第i个元素表示第i个柱子及右边位置中，height的最大高度。
• 数组的初始化：
  
  
• 数组的更新规则：
  当时，
  当时，
• 于是，每个柱子接雨水的数量等于：

官方解法的图片，很清楚地解释了雨水数量计算的原因：

（2）Python实现

class Solution(object):
    def trap(self, height):
        """
        :type height: List[int]
        :rtype: int
        """
        # 动态规划
        # 时间复杂度=O(n)，空间复杂度=O(n)

        # 1. 初始化
        if not height:
            return 0
        n = len(height)
        leftMax = [height[0]] + [0] * (n - 1)
        rightMax = [0] * (n - 1) + [height[-1]]

        # 2. 从左、右扫描
        for i in range(1, n):
            leftMax[i] = max(leftMax[i - 1], height[i])
        for i in range(n - 2, -1, -1):
            rightMax[i] = max(rightMax[i + 1], height[i])

        # 3. 遍历，计算height[i]可以接到的雨水
        ans = 0
        for i in range(n):
            ans += min(leftMax[i], rightMax[i]) - height[i]

        return ans

2、双指针

时间复杂度：O(n)

空间复杂度：O(1)

（1）算法原理

• 首先，维护两个指针left和right、两个变量leftMax和rightMax。
• 初始化：
• 当left指针和right指针没有相遇时，变量的更新规则：

1、使用height[left]和height[right]更新leftMax和rightMax；

2、当height[left]

3、当height[left]>=height[right]时，必有leftMax>=rightMax；此时，第left个柱子能接到的雨水数量等于rightMax-height[left]。

（2）Python实现

class Solution(object):
    def trap(self, height):
        """
        :type height: List[int]
        :rtype: int
        """
        # 双指针
        # 时间复杂度=O(n)，空间复杂度=O(1)

        # 初始化
        if not height:
            return 0
        n = len(height)
        left, right = 0, n - 1
        leftMax, rightMax = 0, 0

        # 移动指针
        ans = 0
        while left < right:
            leftMax = max(leftMax, height[left])
            rightMax = max(rightMax, height[right])
            if height[left] < height[right]:
                ans += leftMax - height[left]
                left += 1
            else:
                ans += rightMax - height[right]
                right -= 1

        return ans

 二、接雨水 II

leetcode.407 链接：https://leetcode.cn/problems/trapping-rain-water-ii/

难度：hard

这道题比较好的解法是用最小堆来维护，三维情况接雨水的最小堆解法是从二维情况延伸的，所以先讲一下二维情况接雨水的最小堆解法~

注：二维情况同接雨水I，为了比较二维和三维情况，所以放在了这里介绍

1、最小堆：二维情况（同接雨水 I）

时间复杂度：O(nlogm)   这里的m指堆中的元素个数，最多为2，可以认为是常数级

空间复杂度：O(n)

（1）算法原理

• 首先，先让两边的端点入堆；同时，使用visited变量记录位置是否被访问。
• 重复以下步骤：
  1、d出堆中较小的元素和位置
  2、检查该元素的左、右高度是否比它高
  3、如果比它低，说明可以注水，则注水到相同高度，并将注水后的高度入堆；如果比它高，直接入堆
• 最终，总的注水量就是接雨水的量

（2）Python实现

class Solution(object):
    def trap(self, height):
        """
        :type height: List[int]
        :rtype: int
        """
        # 最小堆
        # 时间复杂度=O(nlogm)，空间复杂度=O(n)

        if not height:
            return 0
        n = len(height)

        # 初始化：访问了最左和最右
        visited = [True] + [False] * (n - 2) + [True]
        Q = []
        heapq.heappush(Q, (height[0], 0))
        heapq.heappush(Q, (height[-1], n - 1))

        ans = 0
        while Q:
            # d出较小者
            h, pos = heapq.heappop(Q)   
            # 看左右的情况
            for p in [pos + 1, pos - 1]:
                if 0 <= p < n and not visited[p]:
                    if height[p] < h:
                        ans += h - height[p]
                    heapq.heappush(Q, (max(height[p], h), p))
                    visited[p] = True

        return ans

2、最小堆：三维情况

时间复杂度：O(m*n*log(m+n))

空间复杂度：O(m*n)

（1）算法原理

 和二维情况的思路一样，在考虑边界情况和方向的时候，需要考虑到三维。

（2）Python实现

class Solution(object):
    def trapRainWater(self, heightMap):
        """
        :type heightMap: List[List[int]]
        :rtype: int
        """
        m, n = len(heightMap), len(heightMap[0])
        if m <= 2 or n <= 2:
            return 0

        # 将边界入堆
        Q = []
        visited = [[0] * n for _ in range(m)]
        for i in range(m):
            heapq.heappush(Q, (heightMap[i][0], i, 0))
            heapq.heappush(Q, (heightMap[i][n - 1], i, n - 1))
            visited[i][0] = 1
            visited[i][n - 1] = 1
        for j in range(n):
            heapq.heappush(Q, (heightMap[0][j], 0, j))
            heapq.heappush(Q, (heightMap[m - 1][j], m - 1, j))
            visited[0][j] = 1
            visited[m - 1][j] = 1
        
        # 看四周的柱子
        res = 0
        dir = [(-1, 0), (1, 0), (0, -1), (0, 1)]
        while Q:
            height, x, y = heapq.heappop(Q)
            for dx, dy in dir:
                i, j = x + dx, y + dy
                if 0 <= i < m and 0 <= j < n and not visited[i][j]:
                    if heightMap[i][j] < height:
                        res += height - heightMap[i][j]
                    heapq.heappush(Q, (max(heightMap[i][j], height), i, j))
                    visited[i][j] = 1
        
        return res

参考：力扣