LC-805.数组的均值分割

题目描述

leetcode 困难题

给定你一个整数数组 nums

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

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

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

示例1：

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

示例2：

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

提示：

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

分析

设我们移动了 $k$ 个元素到数组 $A$ 中，移动了 $k_2$ 或者说 $n-k$ 个元素到数组 $B$ 中，分割后的数组为 $A$ 和 $B$，且${sum}(A)$， ${sum}(B)$， ${sum}(nums)$ 分别表示数组 $A$， $B$， $nums$ 的元素和，由于数组 $A$，$B$ 的平均值相等，可以推出：

$$
\begin{align}
&\frac{sum(A)}{k} = \frac{sum(B)}{n-k}\\
\Leftrightarrow & sum(A) \times n - sum(A) \times k = sum(B) \times k\\
\Leftrightarrow & sum(A) \times n = (sum(B) + sum(A)) \times k\\
\Leftrightarrow & sum(A) \times n = (sum(nums)) \times k\\
\Leftrightarrow & \frac{sum(A)}{k} = \frac{sum(nums)}{n}\\
\end{align}
$$

也就是说如果存在子数组 $A$，$B$ 的平均值相等，则该平均值必然还会等于 $nums$ 的平均值。

并且我们只需要求出是否存在一个子数组的平均值等于原数组平均值即可，因为如果存在一个子数组平均值等于原数组平均值的话：
$$
\begin{align}
&\frac{sum(nums)}{n} = \frac{sum(A)}{k}\\
\Leftrightarrow & \frac{sum(A) + sum(B)}{k + k_2} = \frac{sum(A)}{k}\\
\Leftrightarrow & \frac{sum(A) + sum(B)}{k_2} = \frac{sum(A) \times (k + k_2)}{k \times k_2}\\
\Leftrightarrow & \frac{sum(B)}{k_2} = \frac{sum(A) \times (k + k_2) - (sum(A) \times k)}{k \times k_2}\\
\Leftrightarrow & \frac{sum(B)}{k_2} = \frac{sum(A) \times k_2 }{k \times k_2}\\
\Leftrightarrow & \frac{sum(B)}{k_2} = \frac{sum(A)}{k}\\
\end{align}
$$
可以得出剩余元素组成的子数组的平均值也为原数组平均值。

那么问题变成求是否存在子数组平均值等于原数组，容易想到的思路是枚举每个子数组，此时一共有 $2^n$ 种方案（每个元素取或不取），由于题目中 n 最大是 30，会 TLE。

我们可以使用折半查找的方法，将时间复杂度降低到 $O(2^\frac{n}{2})$。

我们将数组 $nums$ 分成左右两部分，那么子数组 $A$ 可能存在三种情况：

• 子数组 $A$ 完全在数组 $nums$ 的左半部分；
• 子数组 $A$ 完全在数组 $nums$ 的右半部分；
• 子数组 $A$ 一部分在数组 $nums$ 的左半部分，一部分在数组 $nums$ 的右半部分。

我们分两次来处理数组，每次处理一部分。
如果是前两种情况比较好处理，我们只需要在处理对应部分数组时计算平均值，再与原数组的平均值比较即可。

但如果是第三种情况，则要等我们在处理右半部分时，回过头来通过某种方式搜索左边是否存在符合我们预期的方案（无法简单的通过在哈希表中保存总和来求，因为还和左边数组取了几个有关系），即会花费额外的时间（可能TLE）又难实现。

所以我们可以将数组 ${nums}$ 中的每个元素减去 $nums$ 的平均值，这样数组的平均值则变为 0。那么此时题目中的问题则变为：能否从 $nums$ 中找出若干个元素组成集合 $A$，使得 $A$ 的元素之和为 0。

但是如果直接减去 $\frac{sum(nums)}{n}$ 会引进浮点数，如果在 Java 中直接使用 double 进行运算，会因为浮点数的误差出现WA，所以更好的方案是先将 $nums$ 中的每个元素乘以 $n$ 后再减去数组总和，也就是：
$$
\begin{align}
&nums[i] = nums[i] - \frac{sum(nums)}{n}\\
\Rightarrow &nums[i] = nums[i] \times n - \frac{sum(nums) \times n}{n}\\
\end{align}
$$
这样就不会出现浮点数了。

这样处理后，如果是第一、二种情况，我们只需要在遍历时判断当前方案总和是否等于 $0$ 即可，如果是就直接返回。对于第三种情况，我们也只需要在遍历左边部分时额外再把当前方案的总和保存到哈希表中，然后在遍历右边部分数组的时候，如果当前总和是 $curSum$ ，判断哈希表是否存在 $-curSum$ 即可。

另外需要注意的是，我们不能同时选择左右两边的所有元素，这样数组 $B$ 就为空了。

class Solution {
    public boolean splitArraySameAverage(int[] nums) {
        if(nums.length <= 1){
            return false;
        }
        int n = nums.length, m = nums.length / 2;
        int sum = 0;
        for(int i = 0; i < n; i++){
            sum += nums[i];
        }
        for(int i = 0; i < n; i++){
            nums[i] = nums[i] * n - sum; // 避免浮点数
        }
        Set<Integer> exist = new HashSet<>();
        for(int i = 1; i < (1 << m); i++){ // 2 ^（n/2）个方案
            int curSum = 0;
            for(int j = 0; j < m; j++){
                if((i & (1 << j)) != 0){ // 当前位为 1 代表当前位被选到了，例如 011 代表该轮只选择第一个和第二个元素
                    curSum += nums[j];
                }
            }
            if(curSum == 0){
                return true;
            }
            exist.add(curSum);
        }
        int rSum = 0;
        for (int i = m; i < n; i++) {
            rSum += nums[i];
        }
        for(int i = 1; i < (1 << (n - m)); i++){
            int curSum = 0;
            for(int j = m; j < n; j++){
                if((i & (1 << (j - m))) != 0){
                    curSum += nums[j];
                }
            }
            if(curSum == 0){
                return true;
            }
            if(curSum != rSum && exist.contains(-curSum)){ // 第一个判断是因为不能同时选择左右两边的所有元素
                return true;
            }
        }
        return false;
    }
}