Educational Codeforces Round 23 817-D Imbalanced Array 单调栈

题目链接: Imbalanced Array

大意

定义一个区间不平衡值为: 区间最大值-区间最小值
一个长度为n的数组a, 求其所有子区间的不平衡值的和
$1 \leq n \leq 10^6$

思路

$\sum_{i=0}^{n-1} a[i]*a[i]作为最左最大值出现的区间个数 - a[i]*a[i]作为最左最小值出现的区间个数$
最左最大值: 如果有多个最大值, a[i]是最左边的那个, 保证了不会重复计算
求解区间个数(最大值为例):
l[i]:=从i往左第一个大于等于a[i]的位置
r[i]:=从i往右第一个大于a[i]的位置
那么a[i]作为最左最大值的区间个数 = (r[i]-i)*(i-l[i])

单调栈求解l[i], r[i]:

单调栈: 栈中元素单调, 以单调递增栈为例, 开始时栈内元素都是单调递增的, 将
x入栈是, 如果栈顶元素大于x, 则一直pop直到栈顶元素小于等于x, 再将x入栈
将数组中元素依次入栈, 将大于a[i]的元素全部pop后, 栈顶元素就是左边第一个大于等于a[i]的元素

代码

#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e6+1000;
typedef long long ll;
int n, a[maxn], l[maxn], r[maxn];


int main()
{
    cin >> n;
    for(int i=0; i<n; ++i) scanf("%d", a+i);

    stack<int> s;
    for(int i=0; i<n; ++i) 
    {
        while(!s.empty() && a[i]>a[s.top()]) s.pop();
        l[i] = s.empty() ? -1 : s.top();
        s.push(i);
    }
    while(!s.empty()) s.pop();
    for(int i=n-1; i>=0; --i)
    {
        while(!s.empty() && a[i]>=a[s.top()]) s.pop();
        r[i] = s.empty() ? n : s.top();
        s.push(i);
    }
    ll ans = 0;
    for(int i=0; i<n; ++i) ans += 1ll*a[i]*(r[i]-i)*(i-l[i]);

    while(!s.empty()) s.pop();
    for(int i=0; i<n; ++i) 
    {
        while(!s.empty() && a[i]<a[s.top()]) s.pop();
        l[i] = s.empty() ? -1 : s.top();
        s.push(i);
    }
    while(!s.empty()) s.pop();
    for(int i=n-1; i>=0; --i)
    {
        while(!s.empty() && a[i]<=a[s.top()]) s.pop();
        r[i] = s.empty() ? n : s.top();
        s.push(i);
    }
    for(int i=0; i<n; ++i) ans -= 1ll*a[i]*(r[i]-i)*(i-l[i]);
    cout << ans << endl;

    return 0;
}