题目描述

给定你一个整数数组 nums。

我们要将 nums 数组中的每个元素移动到 A 数组 或者 B 数组中，使得 A 数组和 B 数组不为空，并且 average(A) == average(B)。

如果可以完成则返回 true， 否则返回 false。

注意：对于数组 arr, average(arr) 是 arr 的所有元素除以 arr 长度的和。

样例

输入：nums = [1,2,3,4,5,6,7,8]
输出：true
解释：我们可以将数组分割为 [1,4,5,8] 和 [2,3,6,7], 他们的平均值都是 4.5。

输入：nums = [3,1]
输出：false

提示

• 1 <= nums.length <= 30
• 0 <= nums[i] <= 10^4

算法

(折半枚举) $O(n^2 * 2^{\frac{n}{2}})$

1. 此题要求将数组分为两部分，使得每部分的平均值相等，通过数学推到，不难得知这个平均值也会等于整个数组的平均值。
2. 由于枚举每个数字在哪一部分将会超时，所以我们采取折半枚举的策略，先枚举半个数组中每个数字在哪一部分并记录下属于 A部分的 sum 和数字个数。
3. 然后，再枚举剩下一半数组归属于哪一部分，枚举一组结果后，需要判定是否能和前一半检索过程中出现的结果合并“凑出”平均值。
4. 判定的过程如下，假设后一半数组的一个枚举情况为 (sum, num)， 表示属于 A 部分的总和和数字个数，我们枚举一个系数 c 满足 $\max(1, num) <= c < c$ 并且保证 $c \times avg$ 整数。也就是说， $c \times avg$ 将是目标总和 $t$。然后我们在之前存储的前一半数组的结果中查询，(t - sum, c - num) 是否存在。如果存在，则找到了答案，也就是属于 A 部分的总和为 $t$，数字个数为 $c$，直接返回 true。
5. 如果枚举结束后也没有找到答案，返回 false。

时间复杂度

• 折半枚举的情况共 $O(2^{\frac{n}{2}})$，在前半部分的枚举中更新哈希表需要 $O(n)$ 的时间。（因为 pair 类型不支持 hash 操作，所以只能通过 set 数据结构在对数时间内查询 $O(\log (2^{\frac{n}{2}})) = O(n)$。
• 在后半部分的枚举中，需要 $O(n)$ 的时间枚举 $c$，查询哈希表需要 $O(n)$ 的时间。
• 故总时间复杂度为 $O(n^2 * 2^{\frac{n}{2}})$。

空间复杂度

• set 数据结构需要 $O(2^{\frac{n}{2}})$ 的空间。

C++ 代码

class Solution {
public:
    bool splitArraySameAverage(vector<int>& nums) {
        int n = nums.size();

        int tot = 0;
        for (int i = 0; i < n; i++)
            tot += nums[i];

        int m = n / 2;

        set<pair<int, int>> h;
        for (int s = 0; s < (1 << m); s++) {
            int sum = 0, cnt = 0;

            for (int i = 0; i < m; i++)
                if (s & (1 << i)) {
                    sum += nums[i];
                    cnt++;
                }

            h.insert(make_pair(sum, cnt));
        }

        for (int s = 0; s < (1 << (n - m)); s++) {
            int sum = 0, cnt = 0;

            for (int i = m; i < n; i++)
                if (s & (1 << (i - m))) {
                    sum += nums[i];
                    cnt++;
                }

            for (int c = max(1, cnt); c < n; c++)
                if (c * tot % n == 0) {
                    int t = c * tot / n;

                    pair<int, int> lef(t - sum, c - cnt);
                    if (h.find(lef) != h.end())
                        return true;
                }
        }

        return false;
    }
};