截断数组

#include <bits/stdc++.h>
using namespace std;

const int N = 1e6 + 3;

typedef long long LL;

LL a[N];

void solve(){
    int n; cin >> n;
    for(int i = 1; i <= n; i ++){
        cin >> a[i]; a[i] += a[i - 1];
    }
    if(n < 3 || a[n] % 3 != 0){  //特殊情况
        cout << 0 << endl;
        return ;
    }
    else{
        LL res = a[n] / 3;
        LL cnt = 0, ans = 0;
        for(int i = 1; i + 2 <= n; i ++){
            if(a[i] == res) cnt ++;
            if(a[n] - a[i + 1] == res) ans += cnt;
        }
        cout << ans << endl;
    }
}

int main(){
    solve();
    return 0;
}