【二分】C - Number of Ways题解

C - Number of Ways

题目大意

将一数组分为三部分使得每部分的和相同，问有多少种分法。

暴力思路（TLE)

求前缀和与后缀和找到和为ave的下标，然后二重循环判断是否满足条件。

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int MAXN = 5e5 + 5;
int a[MAXN];
ll sum1[MAXN], sum2[MAXN], ave;
vector<int> pos1,pos2;
int main() {
    int n;
    int ans = 0;
    ave = 0;
    scanf("%d",&n);
    for (int i = 0; i < n; i ++) {
        scanf("%d",&a[i]);
        ave += a[i];
    }
    if (n < 3 || ave % 3) ans = 0;
    else {
        ave /= 3;
        sum1[0] = a[0];
        sum2[n-1] = a[n-1];
        if (sum1[0] == ave) pos1.push_back(0);
        if (sum2[n-1] == ave) pos2.push_back(n-1);
        for (int i = 1; i < n; i ++) {
            sum1[i] = sum1[i - 1] + a[i];
            sum2[n-i-1] = sum2[n-i] + a[n-i-1];
            if (sum1[i] == ave) pos1.push_back(i);
            if (sum2[n-i-1] == ave) pos2.push_back(n-i-1);
        }
        for (int i = 0; i < pos1.size(); i ++) {
            for (int j = 0; j < pos2.size(); j ++) {
                if ((sum1[pos2[j]-1] - sum1[pos1[i]]) == ave && pos2[j]-1 > pos1[i]) ans ++;
            }
        }
    }
    printf("%d\n", ans);
}

但实际上并不需要，因为只要保证前缀和到的下标i是小于后缀和到的下标j-1即一定满足条件，中间剩余值之和肯定也是ave。
故只需要在前缀和满足条件的下标中查找满足此条件的下标的个数即可

AC代码

注意lower_bound(x)函数是不大于x

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int MAXN = 5e5 + 5;
int a[MAXN];
ll sum, ave;
vector<int> pos;
int main() {
    int n;
    ll ans = 0;
    ave = 0;
    scanf("%d",&n);
    for (int i = 0; i < n; i ++) {
        scanf("%d",&a[i]);
        ave += a[i];
    }
    if (n < 3 || ave % 3) ans = 0;
    else {
        ave /= 3;
        for (int i = 0; i < n; i ++) {
            sum += a[i];
            if (sum == ave) pos.push_back(i);
        }
        sum = 0;
        for (int i = n - 1; i >= 0; i --) {
            sum += a[i];
            if (sum == ave) ans += lower_bound(pos.begin(),pos.end(),i-1)-pos.begin();
        }
    }
    printf("%lld\n", ans);
}