Codeforces 466C - Number of Ways Solution

 #include <bits/stdc++.h>
using namespace std;
long long choose(long long n,long long k){
    for(long long i=1;i<k;i++){
        n*=n-i;
    }
    while(k){
        n/=k--;
    }
    return n>0?n:0;
}
int main(){
    ios::sync_with_stdio(0);
    cin.tie(NULL);
    cout.tie(NULL);
    int n;
    cin>>n;
    vector<long long> a(n);
    cin>>a[0];
    for(int i=1;i<n;i++){
        cin>>a[i];
        a[i]+=a[i-1];

    }
    long long sum=a[n-1],sum_1=a[n-1]/3,sum_2=(2*a[n-1])/3;
    if(sum_1+sum_2!=sum)cout<<0;
    else if(sum==0){
        int z=0;
        for(int i=1;i<n-1;i++)
            if(a[i]==0)z++;
        long long ans=choose(z,2);
        if(a[0]==0)ans+=z;
        cout<<ans;
    }
    else{
        int total_2=0,tmp_2=0;
        long long ans=0;
        for(int i=0;i<n;i++){
            if(a[i]==sum_2)total_2++;
        }
        for(int i=0;i<n;i++){
            if(a[i]==sum_1)ans+=total_2 - tmp_2;
            else if(a[i]==sum_2)tmp_2++;
        }
        cout<<ans;
    }    
    return 0;
}