题解 | 喜欢切数组的红

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

int main() {
    int n;
    cin >> n;
    vector<int> vec(n);
    int count_positive = 0;
    for (int i = 0; i < n; ++i) {
        cin >> vec[i];
        if(vec[i] > 0) ++count_positive;
    }

    //判断数组中的正数有没有3个以上 没有直接输出0
    if(count_positive < 3) {
        cout << 0 << endl;
        return 0;
    }

    // 计算前缀和
    vector<int> prefix_sum(n + 1, 0);
    for (int i = 1; i <= n; ++i) {
        prefix_sum[i] = prefix_sum[i - 1] + vec[i - 1];
    }

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

    // 记录 [0, i) 内最后一个正数的下标
    vector<int> max_idx(n + 1, -1);
    for (int i = 1; i <= n; ++i) {
        if (vec[i - 1] > 0) {
            max_idx[i] = i - 1;
        } else {
            max_idx[i] = max_idx[i - 1];
        }
    }

    // 找到所有满足 prefix_sum[i] == target 的位置（左端点）
    vector<int> left;
    for (int i = 1; i <= n; ++i) {
        if (prefix_sum[i] == target) {
            left.push_back(i);
        }
    }

    // 找到所有满足 prefix_sum[j] == 2 * target 的位置（右端点）
    vector<int> right;
    for (int j = 1; j <= n; ++j) {
        if (prefix_sum[j] == 2 * target) {
            right.push_back(j);
        }
    }

    // 统计满足条件的 (i, j) 对数
    int sol = 0;
    int right_ptr = 0; // 指向 right 数组的指针
    for (int i : left) {
        // 找到第一个 j > i 且 prefix_sum[j] == 2 * target
        while (right_ptr < right.size() && right[right_ptr] <= i) {
            right_ptr++;
        }
        // 遍历所有满足 j > i 的 j
        for (int j_idx = right_ptr; j_idx < right.size(); ++j_idx) {
            int j = right[j_idx];
            // 检查 [0, j) 和 [i, j) 和 [j, n) 是否有正数
            if (max_idx[i] >= 0 && max_idx[j] >= i && max_idx[n] >= j) {
                sol++;
            }
        }
    }
    cout << sol << endl;
    return 0;
}