AcWing 3816. 移动元素

原题链接

题目描述

给定一个长度为 $n$ 的正整数数组 $a_1,a_2,…,a_n$ 。

你需要选择其中一个元素，将其移动至数组中的任意位置（也可以留在原位置）。

我们的目标是，在移动元素操作完成以后，将数组分为前后两个非空部分，并使前一部分的各元素之和等于后一部分的各元素之和。

请问，该目标能否达成？

输入格式

第一行包含整数 $T$ ，表示共有 $T$ 组测试数据。

每组数据第一行包含整数 $n$ 。

第二行包含 $n$ 个整数 $a_1,a_2,…,a_n$ 。

输出格式

每组数据输出一行结果，目标可以达成，则输出 YES，否则输出 NO。

数据范围

$1≤T≤20$ ,
$1≤n≤10^5$ ,
$1≤a_i≤10^9$ 。
同一测试点内所有 $n$ 的和不超过 $10^5$ 。

样例

输入样例：

输出样例：

YES
NO
YES

思路

能否达成，可以分三种情况来讨论。

假设 $S_i$ 是前 $i$ 个数的和

存在某一个前缀和 $S_i = \frac{S_n}{2}$ 。不需要移动。
存在某一个后缀和 $S_i = \frac{S_n}{2}$ 。不需要移动。
需要移动的情况下。来看以 $i$ 结尾部分 $a_{1}, a_{2}, \ldots, a_i$ 。如果以 $i$ 结尾的前缀和减去 $a_{k}$ 等于 $\frac{S_n}{2}$ 的话，是不是就相当于给 $a_{k}$ 移动到 $i$ 后面去了，就相当于 $S_i - a_{k} = \frac{S_n}{2}$ ，这里 $S_i$ 和 $S_n$ 都是已知的，整理一下可以得到 $a_{k} = S_i - \frac{S_n}{2}$ 。判断一个数是否存在可以用哈希表来做。

细节：

后缀和就是从后面开始的前缀和。多开一个数组存逆序的数组。
如果总和为奇数肯定不可能。
如果是第 $1$ 种或第 $2$ 种情况，只要在哈希表里插入一个 $0$ 即可。

代码

C++

#include <iostream>
#include <unordered_set>

using namespace std;

typedef long long LL;

const int N = 1e5 + 5;
int n, a[N], b[N];
LL s[N];

bool check(int w[])
{
    for ( int i = 1; i <= n; i ++ ) 
        s[i] = s[i - 1] + w[i];
    if ( s[n] % 2 ) return false;
    unordered_set<LL> S;
    S.insert(0);
    for ( int i = 1; i <= n; i++ )
    {
        S.insert(w[i]);
        if ( S.count(s[i] - s[n] / 2) ) return true;
    }
    return false;
}

int main()
{
    int T;
    cin >> T;
    while ( T -- )
    {
        cin >> n;
        for ( int i = 1, j = n; i <= n; i ++, j -- )
        {
            cin >> a[i];
            b[j] = a[i];
        }
        if ( check(a) || check(b) ) puts("YES");
        else puts("NO");
    }
    return 0;
}

Java

import java.util.*;

public class Main
{
    private static Scanner in = new Scanner(System.in);
    private static final int N = (int) (1e5) + 5;
    private static int n, a[] = new int[N], b[] = new int[N];
    private static long s[] = new long[N];
    
    public static void main(String args[])
    {
        int T = in.nextInt();
        while ( T -- > 0 )
        {
            n = in.nextInt();
            for ( int i = 1, j = n; i <= n; i ++, j -- )
            {
                a[i] = in.nextInt();
                b[j] = a[i];
            }
            System.out.printf("%s\n", check(a) || check(b) ? "YES" : "NO");
        }
    }
    
    private static boolean check(int w[])
    {
        for ( int i = 1; i <= n; i ++ )
            s[i] = s[i - 1] + w[i];
        if ( s[n] % 2 == 1 ) return false;
        Set<Long> S = new HashSet<>();
        S.add(0L);
        for ( int i = 1; i <= n; i ++ )
        {
            S.add((long) w[i]);
            if ( S.contains(s[i] - s[n] / 2) ) return true;
        }
        return false;
    }
}