主题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. 分析题目：

• 从题目可以知道，当后面的柱子高度低于前一个柱子的时候是无法攒到青豆的，比如[3,2,1]
• 当找到一根比前面高的柱子，就可以计算攒到的青豆
• 一旦攒到青豆就可以认为那部分青豆也变成了柱子的一部分
• 从上面三个分析可以很简单的发现只要维护一个单调递减栈就可以计算青豆数量

2. 模拟样例

• 当前栈空，青豆数量为0
• 第一个柱子入栈，栈为[5]，青豆数量为0
• 第二个柱子入栈，栈为[5,0]，青豆数量为0
• 第三个柱子无法入栈，不满足单调递减栈要求，所以将栈顶拿出，此时意味这第一个柱子与第三个柱子之间可以攒青豆，数量为(2-0-1) * (2-0) = 2,此时青豆数量为2。这时可认为第二个柱子高度变为2，将第三个柱子入栈，栈为[5,2]
• 第四个柱子入栈，栈为[5,2,1],青豆数量为2
• 第五个柱子无法入栈，不满足单调递减栈要求，所以将栈顶拿出，此时意味这第五个柱子与第三个柱子之间可以攒青豆，数量为(4-2-1) * (2-1) = 1,此时青豆数量为2+1=3。
• 继续比较，此时栈为[5,2],第五个柱子无法入栈，不满足单调递减栈要求，所以将栈顶拿出,此时意味这第五个柱子与第一个柱子之间可以攒青豆，数量为(4-0-1) * (4-2) = 6,此时青豆数量为3+6 = 9。
• 继续比较，可以入栈，栈为[5,4]
• ...继续模拟最后得到答案为17

3. 计算方法/代码思路

   使用单调栈存储柱子下标(索引)，记栈顶为top，栈顶下一个为left，高度数组为height[n],则有$height[top] < height[left]$

   记此时不满足入栈条件的下标为$i$,则有计算攒豆数量公式为

代码实现

C++实现

#include <iostream>
#include <vector>
#include <stack>
using namespace std;


int Solution(vector<int> const &  vec) {
    
    int length = vec.size();
    int res = 0;
    stack<int> sk;
    for(int i = 0; i<length; i++) {
        while(!sk.empty() && vec[i] > vec[sk.top()]) {
            int top = sk.top();
            sk.pop();
            if(sk.empty()) break;
            int left = sk.top();
            int curWidth = i - left - 1;
            int curHeight = min(vec[left], vec[i]) - vec[top];
            res += curWidth*curHeight; //计算青豆数量
        }
        sk.push(i);
    }
    return res;
}


int main() {
    vector<int> vec = {5,0,2,1,4,0,1,0,3};
    cout << Solution(vec) << endl;
}

Python实现

因为掘金C++代码功能还在测试，没办法运行，所以写了个python版本的

代码链接：主题4——攒青豆

def Solution(vec):
    stk = []
    res = 0
    for index,height in enumerate(vec) :
        while(len(stk)!=0 and height > vec[stk[-1]]):
            top = stk[-1]
            stk.pop()
            if(len(stk)==0):
                break
            left = stk[-1]
            curWidth = index - left - 1
            curHeight = min(vec[left], height) - vec[top]
            res += curWidth*curHeight #计算青豆数量
        stk.append(index)
    return res
def main() -> None:
    vec = [5,0,2,1,4,0,1,0,3]
    print(Solution(vec))
if __name__ == '__main__':
    main()

输出结果