跳转至

42. 接雨水

力扣链接(困难):https://leetcode.cn/problems/trapping-rain-water

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

示例 1:

img

Text Only
1
2
3
输入: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:

Text Only
输入:height = [4,2,0,3,2,5]
输出:9

提示:

  • n == height.length
  • 1 <= n <= 2 * 104
  • 0 <= height[i] <= 105

个人题解

暴力法 \(O(N^2)\)

注意,本代码在力扣最新测试用例中会 TLE。

暴力思路是,首先遍历每一列(注意首列和末列是无法存水的,所以不用遍历),然后从当前遍历列分别向左/向右找到最高的一个柱子,可以理解为找到杯子的杯壁,高度在杯壁以下的位置就是可以存水的位置,用左右侧最高柱子中较低的那个为基准,减去当前位置柱子的高度,即是当前位置的存水量。

本方法其实是双指针,只是时间复杂度并没有得到优化,达到了非常差的 \(O(N^2)\)

C++
class Solution {
public:
    int trap(vector<int>& height) {
        int result = 0;

        if(height.size() < 3) return 0;
        for(int i = 1; i < height.size() - 1; i++) {
            int lMax = 0, rMax = 0;
            // 分别找到左右两侧最高的柱子
            for(int l = i - 1; l >= 0; l--) {
                if(height[l] > lMax) lMax = height[l];
            }
            for(int r = i + 1; r < height.size(); r++) {
                if(height[r] > rMax) rMax = height[r];
            }
            // 计算当前列能接下的水高度
            result += min(lMax, rMax) - height[i] > 0 ? min(lMax, rMax) - height[i] : 0;
        }

        return result;
    }
};

引入动态规划优化双指针算法 \(O(N)\)

不难发现,上述暴力双指针算法有很多重复遍历过程,外层 for 循环每次都会重新对当前元素左右侧的最高高度进行重新计算。所以,对左右两侧最高高度的计算是主要优化点。

然而,代码随想录上对该优化的说明就是一坨屎,轻描淡写来一句:“当前位置左边的最高高度是前一个位置左边最高高度和本高度的最大值”。那这不就是动态规划的范畴吗,这次为什么不进行动态规划五步曲了?并且标题还写着“双指针优化”,别人乍一看还以为使用了双指针来优化了某步骤,实际并非如此。对此,我将标题改为引入动态规划优化双指针算法

所以说,网上的教程看看就好,能不能融会贯通还是得靠自己,具体的分析我写在代码注释中了。

最后,时间复杂度优化到了 \(O(3N) \sim O(N)\),这样就能成功 AC 力扣最新测试用例了。

C++
class Solution {
public:
    int trap(vector<int>& height) {
        if(height.size() < 3) return 0;
        int result = 0;
        vector<int> leftMax(height.size(), 0);
        vector<int> rightMax(height.size(), 0);

        // dp 初始化
        leftMax[0] = height[0];
        rightMax[height.size() - 1] = height[height.size() - 1];
        // 这里是动态规划的思想
        // 当前位置的左/右侧最大高度,来源于 dp[i - 1]
        // 即 dp[i] = max(nums[i], dp[i - 1])
        for(int i = 1; i < height.size(); i++) {
            leftMax[i] = max(height[i], leftMax[i - 1]);
        }
        // 与前面的遍历方向相反,所以状态转移方向也相反
        for(int i = height.size() - 2; i >= 0; i--) {
            rightMax[i] = max(height[i], rightMax[i + 1]);
        }

        for(int i = 1; i < height.size() - 1; i++) {
            int h = min(leftMax[i], rightMax[i]) - height[i];
            if(h) result += h;
        }

        return result;
    }
};

单调栈 \(O(N)\)

单调栈的雨水计算方法和前面两种不同,单调栈是横向计算,每次计算一层或多层组成的矩形面积。

基本思路是,从左到右遍历元素,如果当前元素小于等于栈顶元素,则将当前元素入栈,即保证栈底元素是最大的,和前面两种方法中找左侧最高的柱子是一个逻辑。

关键点:当前元素如果大于栈顶元素,代表此时在当前元素之前出现了凹槽,那么按照每次计算一层或多层组成的矩形面积的逻辑,此时就应该进行计算雨水面积了。

计算方式和前两种方法类似,首先找到左侧第一个大于等于当前元素的位置,只需使用 while 循环,将比当前元素小的栈顶元素依次出栈,这样就找到了左侧第一个大于等于当前元素的元素位置,作为凹槽左边界,这也是单调栈的经典应用。

需要注意的是,由于是单调递减栈,所以每一次 while 循环,当前栈顶元素都会比前一个出栈元素 (mid) 大,所以新的栈顶与当前元素之间的区域会组成一个新的凹槽,mid 作为凹槽深度,计算面积。

这里与前两种方法不同的是,单调栈只能找到某一侧的比当前元素更大/小的位置,不能同时找两侧,所以我们把当前元素看作为右侧的临时最大元素,即右边界。这样也就得到了雨水矩形区域的宽高计算公式:

  • h = min(height[st.top()], height[i]) - height[mid]
  • w = i - st.top() - 1

然后将当前元素入栈(刚才已经用 while 将所有比当前元素小的栈顶出栈,所以当前元素入栈后依旧满足单调递减栈),继续向后遍历,同理处理右方或上方更宽的凹槽,处理顺序如图所示:

image-20250925212448759

为什么将左侧第一个大于等于当前元素的元素位置作为左边界?

因为如果遇到某个元素大于了右边界,其实水坑的最大高度也就是右边界了,继续往左找左边界已经没有意义了,因为水会从右边界漏出去。

C++
class Solution {
public:
    int trap(vector<int>& height) {
        if (height.size() < 3)
            return 0;
        stack<int> st;
        int result = 0;

        st.push(0);
        for (int i = 1; i < height.size(); i++) {
            // 如果当前元素比栈顶元素大,则出现了凹槽
            // 栈顶元素即凹槽对应的下标
            // 如果当前元素 <= 栈顶元素,则将其 push 到栈顶,等待凹槽出现
            if (height[i] <= height[st.top()]) {
                st.push(i);
            } else {
                while (!st.empty() && height[i] > height[st.top()]) {
                    int mid = st.top();
                    st.pop();
                    if (!st.empty()) {
                        // 如果单调栈中有连续一致的元素,其实不影响,因为这个地方 h 会计算为 0
                        // 有个办法可以除掉这个情况,就是把 if (height[i] <= height[st.top()]) 中
                        // 等于的情况独立出来,pop 掉栈顶后在 push,不过不影响结果,只是会影响计算逻辑
                        int h = min(height[st.top()], height[i]) - height[mid];
                        int w = i - st.top() - 1;
                        result += h * w;
                    }
                }
                st.push(i);
            }
        }
        return result;
    }
};

实话说,本题单调栈方法不太好理解,代码随想录讲的又是一坨,所以自己多下功夫吧,形成适合自己的方法论。