主题 4:攒青豆
现有 n 个宽度为 1 的柱子,给出 n 个非负整数依次表示柱子的高度,排列后如下图所示,此时均匀从上空向下撒青豆,计算按此排列的柱子能接住多少青豆。(不考虑边角堆积)
以下为上图例子的解析:
输入:height = [5,0,2,1,4,0,1,0,3]
输出:17
解析:上面是由数组 [5,0,2,1,4,0,1,0,3] 表示的柱子高度,在这种情况下,可以接 17 个单位的青豆。
个人思路
对于能够接住多少青豆来说,首先是要看一个桶的大小,而大小取决于桶的最短边,同样的对于这道题中,我们题目中的桶是一个大的桶,然而我们可以把每一列都看作一个桶,把每一列接到的数量和起来,就是总共的数量了。
现在的问题就转化为求出每一列的桶能够接到多少青豆,这取决于这一列左边与右边的最高列,所以如果找到左边与右边的最高列,再通过最短边可以计算出这一列的数量。例:第四列的高度为 1,对于第四列来说,左边的最高列为第一列 5,右边的最高列为第五列 4,所以说这个桶的最短为 min(5,4)也就为 4,这一列所能够接住的数量也就为(4-1)。
单调栈
那么我们怎么找到左边和右边的最高列呢?这里就可以使用单调栈,如果我们想找左边的最高列,使用单调递增栈,我们就从左向右遍历,元素依次入栈,如果元素大于栈顶元素,就栈顶元素弹出,直到无法弹出时,栈底的元素就为左边的最大值,再将元素入栈。同理如果想求右边的最大值,就从后往前遍历。
代码
func peas(height []int) int {
n := len(height)
l, r := make([]int, n), make([]int, n)
stack := []int{}
ans := 0
for i := 0; i < n; i++ {
for len(stack) > 0 && height[i] >= stack[len(stack)-1] {
stack = stack[:len(stack)-1]
}
if len(stack) > 0 {
l[i] = stack[0]
} else {
l[i] = 0
}
stack = append(stack, height[i])
}
stack = stack[:0]
for i := n - 1; i >= 0; i-- {
for len(stack) > 0 && height[i] >= stack[len(stack)-1] {
stack = stack[:len(stack)-1]
}
if len(stack) > 0 {
r[i] = stack[0]
} else {
r[i] = 0
}
stack = append(stack, height[i])
}
for i, v := range height {
h := min(l[i], r[i])
if h > v {
ans += h - v
}
}
return ans
}
func min(a, b int) int {
if a < b {
return a
}
return b
}