当青训营遇上码上掘金
1. 问题描述
现有 n 个宽度为 1 的柱子，给出 n 个非负整数依次表示柱子的高度，排列后如下图所示，此时均匀从上空向下撒青豆，计算按此排列的柱子能接住多少青豆。（不考虑边角堆积）
以下为上图例子的解析：
输入：height = [5,0,2,1,4,0,1,0,3]
输出：17
解析：上面是由数组 [5,0,2,1,4,0,1,0,3] 表示的柱子高度，在这种情况下，可以接 17 个单位的青豆。
2. 题目分析
- 通过分析可知，可以容纳青豆的关键柱子高度呈先增后减的趋势
- 即对于
height = [2,0,5,0,2,1,4,0,1,0,3]这样一个输入，最终容纳青豆的柱子为[2,5,4,3] - 其余的柱子会被淹没在青豆中，同时会占用青豆的空间
- 可以先通过单调栈来确定容纳青豆的关键柱子
- 然后根据单调栈计算关键柱子围成的空间的容积
- 最后减去非关键柱子所占据的青豆空间
3. 解题思路
- 遍历
height数组，确定height值最高的元素所在位置，如果有多个最大值，随便取一个即可，定义为max_idx - 定义
int类型变量ans保存计算的结果，ans初始化为ans=0 - 从
max_idx开始向左遍历，维护一个height值单调减的索引栈stack，stack中保存的是单调递减的关键柱子的索引 - 当栈
stack空或遍历到的元素i的height值比栈顶元素对应的height值小时，索引i入栈 - 当遍历到的元素
i的height值比栈顶元素对应的height值大时，栈顶元素出栈，出栈的同时ans减去出栈元素对应的height值，表示非关键柱子占据的豆子空间 - 遍历
stack,计算关键柱子所围出的空间，计算结果加到ans中。 - 从
max_idx开始向右遍历，重复第3步至第6步，计算max_idx右侧所能容纳的豆子数，加到ans中 - 返回
ans的值，计算结束
4. 代码实现
def save_green_beans(nums) -> int:
n = len(nums)
max_idx = 0
# 找到值最大的 index，将 index 赋值给 max_idx
for i in range(n):
if nums[i]>nums[max_idx]:
max_idx = i
ans = 0 # 保存结果
# 从 max_idx 开始向左遍历，维护一个单调减的栈，出栈的同时ans减去出栈元素的值，表示比较低的柱子占据的豆子空间
stack = []
for i in range(max_idx, -1, -1):
while len(stack) and nums[stack[-1]]<nums[i]:
ans -= nums[stack[-1]]
stack.pop()
stack.append(i)
# 根据单调栈计算围出的空间，并加到ans
for i in range(len(stack)-1):
ans += nums[stack[i+1]]*(stack[i]-stack[i+1]-1)
# 从 max_idx 开始向右遍历，维护一个单调减的栈，出栈的同时ans减去出栈元素的值，表示比较低的柱子占据的豆子空间
stack = []
for i in range(max_idx, n):
while len(stack) and nums[stack[-1]]<nums[i]:
ans -= nums[stack[-1]]
stack.pop()
stack.append(i)
# 根据单调栈计算围出的空间，并加到ans
for i in range(len(stack)-1):
ans += nums[stack[i+1]]*(stack[i+1]-stack[i]-1)
# 返回结果
return ans
def main() -> None:
nums = [5,0,2,1,4,0,1,0,3]
ans = save_green_beans(nums)
print(ans)
if __name__ == '__main__':
main()
5. 复杂度分析
- 对于height数组中的元素个数N，时间复杂度为O(N)，空间复杂度为O(N)
