题解 | 喜欢切数组的红

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

int main() {
    int n;
    cin >> n;
    vector<int> nums(n + 1);
    vector<int> lt(n + 1); // lt[i] 表示从 1 到 i 的正数个数
    vector<int> sum(n + 1); // sum[i] 表示从 1 到 i 的总和
    vector<int> f(n + 2, 0); // f[i] 表示从 i 到 n 的满足条件的划分数量
    vector<int> nt(n + 2); // nt[i] 表示从 i 开始的第一个正数的位置

    // 输入数组并计算前缀和与正数个数
    for (int i = 1; i <= n; ++i) {
        cin >> nums[i];
        lt[i] = lt[i - 1] + (nums[i] > 0);
        sum[i] = sum[i - 1] + nums[i];
    }

    // 如果总和不能被 3 整除，直接输出 0
    if (sum[n] % 3 != 0) {
        cout << 0 << endl;
        return 0;
    }

    // 初始化 nt 数组
    nt[n + 1] = n + 1;
    for (int i = n; i >= 1; --i) {
        nt[i] = (nums[i] > 0) ? i : nt[i + 1];
    }

    // 初始化 f 数组
    for (int i = n; i >= 1; --i) {
        f[i] = f[i + 1];
        if (lt[n] - lt[i - 1] > 0 && sum[n] - sum[i - 1] == sum[n] / 3) {
            f[i]++;
        }
    }

    // 计算答案
    long long ans = 0;
    for (int i = 1; i <= n; ++i) {
        if (lt[i] > 0 && sum[i] == sum[n] / 3) {
            int r = nt[i + 1];
            if (r < n) {
                ans += f[r + 1];
            }
        }
    }

    // 输出答案
    cout << ans << endl;
    return 0;
}