Codeforces 808D 数组分割

题目大意

给定一个只有正整数组成的数组，判定是否能够只通过一次数字替换来保证能够找到一个合适的分隔位置来使得被分割出来的两半数组，它们分别的和相等。
题目链接见此：Codeforces 808D

这个问题可以抽象成一个确定的前缀数组加上或者去掉一个数字，定义数组：$pre[i]$ 定义一个记录 $[0, i)$ 的数组和，$\sigma$ 表示在可以二分数组和的情况下的二分和，$k$ 表示在 $[0, i)$ 中的一个下标值，那么 $pre[i]$ 和 $\sigma$ 的关系可以表示为：

$$pre_{i} = \begin{cases} \sigma + arr_{i}&\text arr_{i}{移动到左侧},\\ \sigma - arr_{k}&\text arr_{k}{从左侧移出} \end{cases}$$

代码实现如下：

#include <iostream>
#include <cstring>
#include <algorithm>
#include <cstdlib>
#include <cstdio>

using namespace std;
const int maxn = 1e5 + 2;

long long a[maxn];
long long pre[maxn];

bool search(int l, int r, long long ans) {
    while (l < r) {
        int mid = (l + r) / 2;
        if (pre[mid] == ans) return true;
        else if (pre[mid] < ans) l = mid + 1;
        else if (pre[mid] > ans) r = mid;
    }
    return false;
}

int main() {
    int n;
    long long total = 0;
    scanf("%d", &n);
    for (int i = 0; i < n; i++) {
        scanf("%I64d", a + i);
        total += *(a + i);
    }
    a[n] = 0;
    pre[0] = a[0];
    for (int i = 1; i < n; i++) pre[i] = pre[i - 1] + a[i];
    if (total % 2 != 0) printf("NO\n");
    else {
        total /= 2;
        int i = 1;
        bool state = false;
        while (i < n) {
            if (search(0, i, total - a[i])) {
                state = true;
                break;
            }
            i++;
        }
        if (state) {
            printf("YES\n");
        } else {
            i = 0;
            while (i < n) {
                if (search(i, n, total + a[i])) {
                    state = true;
                    break;
                }
                i++;
            }
            if (state) printf("YES\n");
            else printf("NO\n");
        }
    }
    return 0;
}