当青训营遇上码上掘金
题目:攒青豆
现有 n 个宽度为 1 的柱子,给出 n 个非负整数依次表示柱子的高度,排列后如下图所示,此时均匀从上空向下撒青豆,计算按此排列的柱子能接住多少青豆。(不考虑边角堆积)
例子:
输入:`height = [5,0,2,1,4,0,1,0,3]`
输出:`17`
解析:上面是由数组 `[5,0,2,1,4,0,1,0,3]` 表示的柱子高度,在这种情况下,可以接 17 个单位的青豆。
思路:
对于每个槽来说,它能容纳的青豆的最大高度其实就相当于槽两边柱子高度的最小值。
如果我们可以找出这些值,就可以知道答案。
暴力穷举
对于每个柱子,可以使用两个循环分别找到与当前柱子相邻的两个最高的柱子。
然后用与当前柱子相邻的最高柱子中较低的那个,减去当前柱子的高度,即可得到这个槽里被豆子的数量。
把所有的数量相加,即可得到总数。
def trap(height):
res = 0
for i in range(len(height)):
max_left = 0
max_right = 0
for j in range(i, -1, -1):
max_left = max(max_left, height[j])
for j in range(i, len(height)):
max_right = max(max_right, height[j])
res += min(max_left, max_right) - height[i]
return res
height = [5,0,2,1,4,0,1,0,3] #测试输入
print("青豆数量:", trap(height))
因为需要在两个内部循环中枚举每个柱子并找到与它相邻的两个最高的柱子,并且每个柱子都需要在外部循环中枚举,所以时间复杂度为 O(n^2)。
栈维护法
在暴力穷举的过程中,有一些比较矮的柱子被多次计算了,这很没有必要。
所以我们可以用栈来保存当前遇到的最高的柱子。
基本思路是:
- 初始化一个空栈。
2. 从左到右遍历数组。
3. 当前柱子高度大于栈顶元素时,将栈顶元素弹出并计算栈顶元素与当前柱子高度之间的积水。
4. 重复步骤 3,直到栈顶元素比当前柱子高度小或相等,将当前柱子入栈。
5. 重复步骤 2-4,直到遍历完整个数组。
def trap(height):
if not height:
return 0
n = len(height)
stack = []
res = 0
for i in range(n):
while stack and height[i] > height[stack[-1]]:
top = stack.pop()
if not stack:
break
distance = i - stack[-1] - 1
bounded_height = min(height[i], height[stack[-1]]) - height[top]
res += distance * bounded_height
stack.append(i)
return res
height = [5,0,2,1,4,0,1,0,3] #测试输入
print("青豆数量:", trap(height))
因为每个柱子最多入栈和出栈一次,所以栈法的时间复杂度是 O(n)。