说真的,接雨水这道题,我敢说刷题的朋友都绕不开。 我第一次碰它的时候,上来就写暴力解法,小例子跑得溜溜的。 一提交直接超时,盯着 O (n²) 的复杂度愣了半天。 后来才反应过来,好多计算都是重复的,完全可以提前算好存起来。 今天就跟大家聊聊这个最容易理解的优化思路 ------ 前后缀最大值数组,看完你也能秒写出来。
题目说的啥
题目很直白。 给你一个数组,每个数字代表一根宽度为 1 的柱子高度。 问下雨之后,这些柱子之间的凹槽,总共能接住多少雨水。 本质就是算出每个位置能存多少水,最后加总就行。
从暴力超时想到的优化思路
我最开始的暴力思路其实很朴素。 每个位置能存多少水,不就是看左边最高的柱子和右边最高的柱子,取更矮的那个,再减去当前柱子的高度嘛。 存不了水就按 0 算。
但问题也出在这。 每算一个位置,都要往左扫一遍找最大值,往右扫一遍找最大值。 数据量一大,时间直接就爆了。
当时我就想,这些最大值其实是有递推关系的啊。 比如第 i 个位置左边的最大值,不就是「第 i-1 个位置的左边最大值」和「当前柱子高度」里更大的那个吗? 根本不用每次都从头扫。
说白了就是空间换时间。 我先花一次遍历的功夫,把每个位置左边的最高高度都算出来,存成一个前缀数组。 再倒着遍历一次,把每个位置右边的最高高度也存成一个后缀数组。 最后再扫一遍,挨个算每个位置的存水量,加起来就完事了。
手算一遍就懂了
拿示例 [4,2,0,3,2,5] 来走一遍,一眼就清晰。
先算左边最大值数组 preMax,从左往右推:
- 第 0 位左边只有自己:4
- 第 1 位:max (前一位最大值 4, 当前高度 2) = 4
- 第 2 位:max (4, 0) = 4
- 第 3 位:max (4, 3) = 4
- 第 4 位:max (4, 2) = 4
- 第 5 位:max (4, 5) = 5 最终 preMax 就是
[4,4,4,4,4,5]
再算右边最大值数组 sufMax,从右往左推:
- 第 5 位右边只有自己:5
- 第 4 位:max (当前高度 2, 后一位最大值 5) = 5
- 第 3 位:max (3, 5) = 5
- 第 2 位:max (0, 5) = 5
- 第 1 位:max (2, 5) = 5
- 第 0 位:max (4, 5) = 5 最终 sufMax 就是
[5,5,5,5,5,5]
最后每个位置取两者的最小值,减去当前柱子高度: 第 0 位:min (4,5)-4 = 0 第 1 位:min (4,5)-2 = 2 第 2 位:min (4,5)-0 = 4 第 3 位:min (4,5)-3 = 1 第 4 位:min (4,5)-2 = 2 第 5 位:min (5,5)-5 = 0 加起来 0+2+4+1+2+0 = 9,和示例答案完全对上。
是不是手算一遍就特别清楚? 其实就是把暴力法里重复的计算,提前批量算好了。
代码实现
思路通了,代码写起来就很顺。 这是整理好的 JavaScript 版本,直接提交就能过:
JavaScript
var trap = function(height) {
const n = height.length;
// preMax[i]:第i个位置(含i本身)左边的最大高度
const preMax = Array(n);
preMax[0] = height[0]; // 第一个位置左边只有自己,别漏了初始化
for (let i = 1; i < n; i++) {
preMax[i] = Math.max(preMax[i - 1], height[i]);
}
// sufMax[i]:第i个位置(含i本身)右边的最大高度
const sufMax = Array(n);
sufMax[n - 1] = height[n - 1]; // 最后一个位置同理
for (let i = n - 2; i >= 0; i--) {
// 别手滑写成i++,我一开始写反直接数组越界,踩过的坑
sufMax[i] = Math.max(sufMax[i + 1], height[i]);
}
let ans = 0;
for (let i = 0; i < n; i++) {
// 左右挡板取更矮的,减去当前柱子高度,就是能存的水量
ans += Math.min(preMax[i], sufMax[i]) - height[i];
}
return ans;
};
复杂度分析
时间复杂度是 O (n),总共遍历了三次数组,每一趟都是线性的。 空间复杂度是 O (n),用了两个长度为 n 的数组来存储最大值。 比起暴力法的 O (n²) 时间,已经提升了一大截,在题目数据范围内完全够用。
我踩过的两个小坑
说两个自己写的时候踩过的坑,大家注意避开。 第一个是边界初始化忘了写。preMax 第一个元素和 sufMax 最后一个元素一定要先赋值,不然循环里取到 undefined,结果直接变 NaN。 第二个是后缀数组的遍历方向搞反。从后往前遍历起始下标是 n-2,步长是 i--,写顺了很容易手滑写成 i++,跑起来直接报错。
最后聊聊
其实这个解法本质就是动态规划的思路,利用之前的计算结果推导出当前值,避免重复劳动。 对于刚刷这道题的朋友来说,这个思路最好理解,也最不容易写错。
你第一次写接雨水用的是什么方法? 是暴力直接莽,还是直接就想到了前后缀优化? 评论区聊聊你的刷题经历,我每条都会看~ 如果觉得文章对你有帮助,点个赞让更多朋友看到呀~