AcWing 4738. 快乐子数组

AcWing 4738. 快乐子数组

我们将 F(B,L,R) 定义为整数数组 B 的索引从 L 到 R（包括两者）的子数组的各个元素之和。

更具体的说，F(B,L,R)=BL+BL+1+…+BR。

如果一个长度为 K 的整数数组 C 满足其所有前缀和均为非负整数，则称数组 C 为快乐数组。

更具体的说，如果 F(C,1,1),F(C,1,2),…,F(C,1,K) 均为非负整数，则数组 C 为快乐数组。

给定一个包含 N 个整数的数组 A，请你计算数组 A 中的所有快乐连续子数组的元素和相加的结果。

输入格式

第一行包含整数 T，表示共有 T 组测试数据。

每组数据第一行包含整数 N。

第二行包含 N 个整数 A1,A2,…,AN。

输出格式

每组数据输出一个结果，每个结果占一行。

结果表示为 Case #x: y，其中 x 为组别编号（从 1 开始），y 为所有快乐连续子数组的元素和相加的结果。

数据范围

1≤T≤100,
−800≤Ai≤800,
每个测试点最多 3030 组数据满足 1≤N≤4×$10^5$，其余数据满足 1≤N≤200。

输入样例：

2
5
1 -2 3 -2 4
3
1 0 3

输出样例：

Case #1: 14
Case #2: 12

样例解释

在 Case 1 中，满足条件的快乐连续子数组有 [1],[3],[3,−2],[3,−2,4],[4],它们的元素和分别为 1,3,1,5,4,，相加得到结果 14。

在 Case 2 中，满足条件的快乐连续子数组有 [1],[1,0],[1,0,3],[0],[0,3],[3],它们的元素和分别为 1,1,4,0,3,3,相加得到结果 12。

ac代码

#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long LL;
const int N = 400010, INF = 1e9;

int n;
int stk[N];
LL A[N], B[N];

int main(){
    int T;
    scanf("%d", &T);

    for (int cases = 1; cases <= T; cases ++ ){
        scanf("%d", &n);
        for (int i = 1; i <= n; i ++ ){
            int x;
            scanf("%d", &x);
            A[i] = A[i - 1] + x;
            B[i] = B[i - 1] + x * (n - i);
        }
        LL res = 0;
        int top = 0;
        A[n + 1] = -INF;
        stk[ ++ top] = n + 1;  // 哨兵
        for (int i = n; i >= 0; i -- ){
            while (A[stk[top]] >= A[i]) top -- ;
            int j = stk[top];
            stk[ ++ top] = i;
            res += B[j - 1] - B[i] - (A[j - 1] - A[i]) * (n - j);
        }
        printf("Case #%d: %lld\n", cases, res);
    }
    return 0;
}