[单调栈]:484.股票价格上涨天数计算

股票价格上涨天数计算

问题描述

小C是一名股票交易员，最近他关注某只股票的价格波动。给定该股票连续N天的价格列表 stockPrices，你需要为小C生成一个新列表，每个位置的值表示从那天起至少需要等待多少天才能看到价格上涨。如果没有上涨的情况，则对应位置的值为0。

例如，对于股票价格列表 [33, 34, 14, 12, 16]，从第一天价格 33 开始，价格上涨发生在第二天价格 34，所以输出 1。若某天之后不再有价格上涨，则输出 0。

测试样例

样例1：

输入：N = 5, stockPrices = [33, 34, 14, 12, 16]
输出：[1, 0, 2, 1, 0]

样例2：

输入：N = 6, stockPrices = [45, 44, 46, 43, 42, 48]
输出：[2, 1, 3, 2, 1, 0]

样例3：

输入：N = 3, stockPrices = [10, 9, 8]
输出：[0, 0, 0]

解题思路

• 朴素做法：遍历每一个数，依次找到后面第一个大于的数，$O(n^2)$
• 考虑优化：
  • 假设数组为$a = [3,2,1,4]$， 可以看出对于前三个数$[3,2,1]$，他们的答案其实都是同一个$[4]$
  • 考虑$a = [3,1,3,4]$， 对于$[1]$来说$[3]$就是答案， 而两个$3$的答案是最后一个$4$
  • 我们可以利用一个存储单调递减的结构，当最小值$mn$小于了当前$a[i]$， 说明对于此时的$mn$答案就是$a[i]$， 然后把这个最小数移除，再次判断，直到最小值大于$a[i]$
  • 可以构建单调递减栈实现
  • 以上面$a = [3,2,1,4]$，我们对于前三个都会放入栈中，
  • 遇到$4$时，栈顶元素为$1$, 更新$ans[栈顶元素下标]$为$i-栈顶元素的下标$，弹出元素$1$
  • 此时栈顶为$2$，...
  • 同理

核心代码(go/py3/c++)

func solution(N int, stockPrices []int) []int {
    ans := make([]int, N)
    st := []int{}
    for i, v := range stockPrices {
        for len(st) > 0 && stockPrices[st[len(st)-1]] < v {
            idx := st[len(st)-1]
            ans[idx] = i - idx 
            st = st[:len(st)-1]
        }
        st = append(st, i)
    }
    return ans
}

def solution(N, stockPrices):
    ans = [0] * N
    st = []
    for i, v in enumerate(stockPrices):
        while st and stockPrices[st[-1]] < v:
            idx = st.pop()
            ans[idx] = i - idx
        st.append(i)
    return ans

vector<int> solution(int N, vector<int>& stockPrices) {
    // write code here
    vector<int>ans(N, 0);
    stack<int>st;
    for(int i=0; i < N; i++) {
        while(!st.empty() && stockPrices[st.top()] < stockPrices[i]) {
            ans[st.top()] = i - st.top();
            st.pop();
        } 
        st.push(i);
    } 
    return ans;
}