[CF466C] Number of Ways

Description

给定一个数组 (a)，包含 (n) 个数字 (a[1],a[2],a[3],...,a[n])。现在你要找到把它分成三份的方法，使得每一份之内所有数的和相等。

Solution

设总和为 (sum)，只考虑 (3|sum) 的情况，设 (k=sum/3)，则统计 (s[i]=sum_{j=1}^{i} [a[i]=k])，则 (ans=sum_{a[i]=2k} s[i-1])。

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

#define int long long 
const int N = 1000005;

signed main()
{
    int n;

    cin>>n;
    vector <int> a(n+2);
    for(int i=1;i<=n;i++) cin>>a[i];

    int sum=0;
    for(int i=1;i<=n;i++) sum+=a[i];

    int k=sum/3;

    int ans=0;

    if(k*3==sum)
    {
        for(int i=1;i<=n;i++) a[i]+=a[i-1];

        vector <int> s(n+2);

        for(int i=1;i<=n;i++) s[i]=s[i-1]+(a[i]==k);

        // for(int i=1;i<=n;i++) cout<<s[i]<<",";
        // cout<<endl;

        for(int i=2;i<n;i++) 
            if(a[i]==2*k) ans+=s[i-1]; 
    }

    cout<<ans<<endl;
    return 0;
}