Codeforces - 466C - Number of Ways - 组合数学

https://codeforces.com/problemset/problem/466/C

要把数据分为均等的非空的三组，那么每次确定第二个分割点的时候把（除此之外的）第一个分割点的数目加上就可以了。记得最后给第三组留至少一个。

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

int n;
int a[500005];

int main(){
    scanf("%d",&n);
    for(int i=0;i<n;i++)
        scanf("%d",&a[i]);

    ll sum=0;
    for(int i=0;i<n;i++){
        sum+=a[i];
    }

    if(sum%3!=0){
        puts("0");
        return 0;
    }

    ll d=sum/3;
    int z1=0;

    ll cnt=0;

    ll cur=0;
    int inc=0;

    int i=0;
    while(i<n){
        cur+=a[i];
        if(cur==d){
            z1++;
            inc=1;
        }
        else{
            inc=0;
        }

        if(z1>=1){
            if(cur==2ll*d){
                if(i<=n-2){
                    //至少把最后一个n-1留给第三组
                    if(inc)
                        cnt+=z1-1;
                    else
                        cnt+=z1;
                    //printf("cnt=%lld i=%d\n",cnt,i);
                }
            }
        }
        i++;
    }

    printf("%lld\n",cnt);
}