[42][困难][栈][动态规划] 接雨水

题目描述

42. 接雨水

给定 n 个非负整数表示每个宽度为 1 的柱子的高度图，计算按此排列的柱子，下雨之后能接多少雨水。

示例 1：

输入：height = [0,1,0,2,1,0,1,3,2,1,2,1]
输出：6
解释：上面是由数组 [0,1,0,2,1,0,1,3,2,1,2,1] 表示的高度图，在这种情况下，可以接 6 个单位的雨水（蓝色部分表示雨水）。

示例 2：

输入：height = [4,2,0,3,2,5]
输出：9

提示：

• n == height.length

• 0 <= n <= 3 * 104

• 0 <= height[i] <= 105

解题思路

按列求

可以看到每一列雨水的高度, 是它左右两侧最高的柱子中, 较矮的那根柱子的高度, 再减去这一列柱子自身的高度.

以[0,1,0,2,1,0,1,3,2,1,2,1]为例, 位置4柱子自身的高度为1, 它左右两侧最高的柱子分别为2和3, 因此对应的水柱的高度为$\min(2, 3) - 1 = 1$.

再如下面这种情况, 中间柱子高度为2, 它左右两侧最高的柱子高度分别为2和3, 它们之中的最小值并不大于当前柱子, 因此对于当前柱子, 是接不到雨水的.

一种暴力的算法是对每一列向左右发散, 找到它对应的左右两侧最高的柱子, 然后根据上述逻辑计算该柱子对应水柱的高度, 再将所有柱子对应的水柱加和在一起即可. 时间复杂度为$O(n^2)$

动态规划

上面的暴力方法中有大量的重复计算, 其实找到每根柱子左右两侧最高的柱子不需要遍历一遍, 使用动态规划的方法计算可以重复利用之前的结果.

先看左侧最高的高度. 使用dp_left记录每个位置其左侧最高柱子的高度. 对于位置i, 相对于i-1, 其左侧多了i-1这个位置的柱子, 因此将dp_left[i-1]与i-1这个位置的新柱子相比, 取最大的即可.

右侧最高的高度也同理求取, 只不过需要从右向左遍历.

class Solution:
    def trap(self, height: List[int]) -> int:
        n = len(height)
        dp_right = [0] * n
        for i in range(n - 2, -1, -1):
            dp_right[i] = max(height[i + 1], dp_right[i + 1])

        max_left, total = 0, 0
        for i in range(1, n):
            max_left = max(max_left, height[i - 1])
            upper = min(max_left, dp_right[i])
            if height[i] < upper:
                total += upper - height[i]

        return total

上面的代码里只有右侧最高柱子使用了数组记录, 左侧最高柱子跟随着迭代更新, 节省空间.

单调栈

以上是按列求的思路. 使用单调栈是按行求取的思路.

从上图中可以看到, 如果我们在右侧发现了比左侧更高的元素, 就可以围成凹槽. 这里的思路不再是求每个柱子对应的接水高度, 而是根据凹槽的深度和宽度, 计算凹槽的容量. 这里的凹槽指的是底部平坦的凹槽. 上图中中间容量为4的部分, 其实是两个凹槽组合而成.

因此我们需要使用一个单调递减栈, 从左向右遍历:

• 当前柱子如果比栈顶元素大, 且栈中至少包含两个元素, 由于栈是递减栈, 说明栈顶元素就是凹槽的底部, 且这个凹槽的左右两侧也确定了

  

  将栈内所有小于当前柱子高度的柱子都弹出, 计算它们分别对应的凹槽深度. 我们记当前柱子为C, 栈顶柱子为B, 如果C的高度大于B的高度, 将B弹出, B的高度是凹槽的底部高度, 凹槽的顶部取它左右两侧的最小值, 即当前柱子C以及弹出B后当前的栈顶柱子A对应的高度. 而凹槽的宽度就是A和C之间的宽度.

  因此栈中记录的不是高度, 而是柱子对应的索引.

• 如果当前柱子与栈顶高度相同, 则它们不可能组成直接组成一个凹槽:

  • 它们可以属于凹槽, 但右侧也有比它们更大的柱子. 因此从这个角度来说栈中存两个之中任意一个都可以, 因为在计算它们代表的凹槽时只提供底部高度, 不提供宽度信息

  • 它们可以是右侧其他柱子的左边边界, 在这种情况下它们提供左边边界索引信息. 而且与这两个相同高度柱子中右侧的组成凹槽

  • 因此当前柱子与栈顶相同时, 直接pop出栈

  

总结过程如下

• 先将下标0入栈

• 遍历之后的每个位置, 如果这个位置的高度小于栈顶元素的高度, 直接入栈

• 如果等于栈顶元素的高度, 将栈顶元素pop出去后入栈

• 如果大于栈顶元素高度, 将栈中所有小于或等于这个高度的元素都pop出去, 并分别计算各个凹槽的容量

  • 取出的栈顶元素是凹槽的底部高度h = height[stack.pop()]

  • 当前柱子高度与栈顶元素弹出之后, 新的栈顶元素的高度中的更小值, 是凹槽的高度

  • 因此对应的凹槽深度为min(height[stack[-1]], height[i]) - h

  • 凹槽对应的宽度为i - stack[-1] - 1

  • 凹槽的容量为两面两个值的乘积

class Solution:
    def trap(self, height: List[int]) -> int:
        n = len(height)
        if n < 3:
            return 0

        stack = [0]
        count  = 0
        for i in range(1, n):
            if height[i] == height[stack[-1]]:
                stack.pop()
            else:
                while stack and height[i] >= height[stack[-1]]:
                    bottom = height[stack.pop()]
                    if stack:
                        upper = min(height[stack[-1]], height[i])
                        width = i - stack[-1] - 1
                        count += (upper - bottom) * width
            stack.append(i)

        return count

参考资料

• 详细通俗的思路分析，多解法
• 42. 接雨水:【双指针】【动态规划】【单调栈】详解！