接雨水

标签: 数组 双指针 动态规划 单调栈

难度: Hard

给定 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
  • 1 <= n <= 2 * 104
  • 0 <= height[i] <= 105

Submission

运行时间: 32 ms

内存: 15 MB

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

        ans = 0
        stack = []
        for i in range(n):
            while stack and height[stack[-1]] < height[i]:
                idx = stack.pop()

                while stack and height[stack[-1]] == height[idx]:
                    stack.pop()

                if stack:
                    left = height[stack[-1]]
                    cur = height[idx]
                    right = height[i]
                    ans += (min(left, right) - cur) * (i-stack[-1]-1)
            stack.append(i)
        return ans

Explain

这个题解使用了单调栈的思路。维护一个单调递减的栈,栈中存储柱子的索引。当遇到一个高度大于栈顶索引对应高度的柱子时,说明形成了一个凹槽,可以接雨水。通过计算凹槽的宽度和高度,可以得到该凹槽能接的雨水量。将这些雨水量累加起来就是最终答案。

时间复杂度: O(n)

空间复杂度: O(n)

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

        ans = 0
        stack = []
        for i in range(n):
            # 当右边的柱子高度大于栈顶柱子的高度,说明形成了凹槽
            while stack and height[stack[-1]] < height[i]:
                idx = stack.pop()
                
                # 将栈中等高的柱子全部弹出
                while stack and height[stack[-1]] == height[idx]:
                    stack.pop()
                
                # 计算凹槽的雨水量
                if stack:
                    left = height[stack[-1]]
                    cur = height[idx]
                    right = height[i]
                    ans += (min(left, right) - cur) * (i-stack[-1]-1)
            stack.append(i)
        return ans

Explore

单调栈的选择是基于其能有效处理问题中涉及的序列局部最大或最小值的特性。在接雨水问题中,我们需要找到每个柱子左右两边比它高的柱子,以确定这个柱子可以接的雨水量。单调栈可以在遍历数组的同时,通过维护一个递减的栈来快速找到左边和右边第一个比当前柱子高的柱子,从而计算雨水量。这种方法相较于数组和队列,可以更直接地解决问题并且减少不必要的重复计算,提高算法效率。

在单调栈中,当遇到一个当前柱子的高度大于栈顶柱子的高度时,栈顶的柱子索引表示凹槽的底部。此时,栈顶下一个元素的索引就是凹槽左边界的柱子,当前索引是凹槽右边界的柱子。凹槽的宽度可以通过右边界的索引减去左边界索引再减一来计算,即 `i - stack[-1] - 1`,这里的 `i` 是当前柱子的索引,`stack[-1]` 是凹槽左边的柱子索引。

在处理等高的柱子时,将栈中等高的柱子全部弹出是为了避免重复计算雨水的存储量。如果栈中有多个等高的柱子,这些柱子之间不可能存储任何雨水,因为它们的高度相同,水会从两者之间流走。只有最左边的等高柱子可能与更左边的一个更高的柱子形成凹槽来存储雨水。因此,保留最左侧的等高柱子,删除其他等高柱子可以简化问题的处理,减少不必要的计算。

在本题解的逻辑中,如果栈顶柱子的高度等于当前柱子的高度,这种情况实际上并没有进行特别的处理。这是因为等高的栈顶柱子和当前柱子之间无法形成有效的凹槽来存储雨水,因此可以视为不影响总雨水量的计算。在实际实现中,遇到此种情况通常意味着当前柱子不会对栈产生影响,或者可以将栈顶的柱子替换为当前柱子的索引,以便后续可能的更高柱子能正确计算宽度。