[LeetCode] 805. Split Array With Same Average 分割数组成相同平均值的小数组

In a given integer array A, we must move every element of A to either list B or list C. (B and C initially start empty.)

Return true if and only if after such a move, it is possible that the average value of B is equal to the average value of C, and B and C are both non-empty.

Example :
Input: 
[1,2,3,4,5,6,7,8]
Output: true
Explanation: We can split the array into [1,4,5,8] and [2,3,6,7], and both of them have the average of 4.5.

Note:

The length of A will be in the range [1, 30].
A[i] will be in the range of [0, 10000].

这道题给了一个数组A，问是否可以把数组分割成两个小数组，并且要求分成的这两个数组的平均值相同。之前我们有做过分成和相同的两个小数组 Split Array with Equal Sum，看了题目中的给的例子，你可能会有种错觉，觉得这两个问题是一样的，因为题目中分成的两个小数组的长度是一样的，那么平均值相同的话，和一定也是相同的。但其实是不对的，很简单的一个例子，比如数组 [2, 2, 2]，可以分成平均值相同的两个数组 [2, 2] 和 [2]，但是无法分成和相同的两个数组。现在唯一知道的就是这两个数组的平均值相等，这里有个隐含条件，就是整个数组的平均值也和这两个数组的平均值相等，这个不用多说了吧，加法的结合律的应用啊。由于平均值是由数字总和除以个数得来的，那么假设整个数组有n个数组，数字总和为 sum，分成的其中一个数组有k个，假设其数字和为 sum1，那么另一个数组就有 n-k 个，假设其数组和为 sum2，就有如下等式：

sum / n = sum1 / k = sum2 / (n - k)

看前两个等式，sum / n = sum1 / k，可以变个形，sum * k / n = sum1，那么由于数字都是整数，所以 sum * k 一定可以整除 n，这个可能当作一个快速的判断条件。下面来考虑k的取值范围，由于要分成两个数组，可以始终将k当作其中较短的那个数组，则k的取值范围就是 [1, n/2]，就是说，如果在这个范围内的k，没有满足的 sum * k % n == 0 的话，那么可以直接返回false，这是一个快速的剪枝过程。如果有的话，也不能立即说可以分成满足题意的两个小数组，最简单的例子，比如数组 [1, 3]，当k=1时，sum * k % n == 0 成立，但明显不能分成两个平均值相等的数组。所以还需要进一步检测，当找到满足的 sum * k % n == 0 的k了时候，其实可以直接算出 sum1，通过 sum * k / n，那么就知道较短的数组的数字之和，只要能在原数组中数组找到任意k个数字，使其和为 sum1，就可以 split 成功了。

问题到这里就转化为了如何在数组中找到任意k个数字，使其和为一个给定值。有点像 Combination Sum III 那道题，当然可以根据不同的k值，都分别找原数组中找一遍，但想更高效一点，因为毕竟k的范围是固定的，可以事先任意选数组中k个数字，将其所有可能出现的数字和保存下来，最后再查找。那么为了去重复跟快速查找，可以使用 HashSet 来保存数字和，可以建立 n/2 + 1 个 HashSet，多加1是为了不做数组下标的转换，并且防止越界，因为在累加的过程中，计算k的时候，需要用到 k-1 的情况。讲到这里，你会不会一拍大腿，吼道，这尼玛不就是动态规划 Dynamic Programming 么。恭喜你骚年，没错，这里的 dp 数组就是一个包含 HashSet 的数组，其中 dp[i] 表示数组中任选 i 个数字，所有可能的数字和。首先在 dp[0] 中加入一个0，这个是为了防止越界。更新 dp[i] 的思路是，对于 dp[i-1] 中的每个数字，都加上一个新的数字，所以最外层的 for 循环是遍历原数组的中的每个数字的，中间的 for 循环是遍历k的，从 n/2 遍历到1，然后最内层的 for 循环是遍历 dp[i-1] 中的所有数组，加上最外层遍历到的数字，并存入 dp[i] 即可。整个 dp 数组更新好了之后，下面就是验证的环节了，对于每个k值，验证若 sum * k / n == 0 成立，并且 sum * i / n 在 dp[i] 中存在，则返回 true。最后都没有成立的话，返回 false，参见代码如下：

解法一：

class Solution {
public:
    bool splitArraySameAverage(vector<int>& A) {
        int n = A.size(), m = n / 2, sum = accumulate(A.begin(), A.end(), 0);
        bool possible = false;
        for (int i = 1; i <= m && !possible; ++i) {
            if (sum * i % n == 0) possible = true;
        }
        if (!possible) return false;
        vector<unordered_set<int>> dp(m + 1);
        dp[0].insert(0);
        for (int num : A) {
            for (int i = m; i >= 1; --i) {
                for (auto a : dp[i - 1]) {
                    dp[i].insert(a + num);
                }
            }
        }
        for (int i = 1; i <= m; ++i) {
            if (sum * i % n == 0 && dp[i].count(sum * i / n)) return true;
        }
        return false;
    }
};

下面这种解法跟上面的解法十分的相似，唯一的不同就是使用了 bitset 这个数据结构，在之前那道 Partition Equal Subset Sum 的解法二中，也使用了 biset，了解了其使用方法后，就会发现使用这里使用它只是单纯的为了炫技而已。由于 biset 不能动态变换大小，所以初始化的时候就要确定，题目中限定了数组中最多 30 个数字，每个数字最大 10000，那么就初始化 n/2+1 个 biset，每个大小为 300001 即可。然后每个都初始化个1进去，之后更新的操作，就是把 bits[i-1] 左移 num 个，然后或到 bits[i] 即可，最后查找的时候，有点像二维数组的查找方式一样，直接两个中括号坐标定位即可，参见代码如下：

解法二：

class Solution {
public:
    bool splitArraySameAverage(vector<int>& A) {
        int n = A.size(), m = n / 2, sum = accumulate(A.begin(), A.end(), 0);
        bool possible = false;
        for (int i = 1; i <= m && !possible; ++i) {
            if (sum * i % n == 0) possible = true;
        }
        if (!possible) return false;
        bitset<300001> bits[m + 1] = {1};
        for (int num : A) {
            for (int i = m; i >= 1; --i) {
                bits[i] |= bits[i - 1] << num;
            }
        }
        for (int i = 1; i <= m; ++i) {
            if (sum * i % n == 0 && bits[i][sum * i / n]) return true;
        }
        return false;
    }
};

再来看一种递归的写法，说实话在博主看来，一般不使用记忆数组的递归解法，等同于暴力破解，基本很难通过 OJ，除非你进行了大量的剪枝优化处理。这里就是这种情况，首先还是常规的k值快速扫描一遍，确保可能存在解。然后给数组排了序，然后对于满足 sum * k % n == 0 的k值，进行了递归函数的进一步检测。需要传入当前剩余数字和，剩余个数，以及在原数组中的遍历位置，如果当前数字剩余个数为0了，说明已经取完了k个数字了，那么如果剩余数字和为0了，则说明成功的找到了k个和为 sum * k / n 的数字，返回 ture，否则 false。然后看若当前要加入的数字大于当前的平均值，则直接返回 false，因为已经给原数组排过序了，之后的数字只会越来越大，一旦超过了平均值，就不可能再降下来了，这是一个相当重要的剪枝，估计能过 OJ 全靠它。之后开始从 start 开始遍历，当前遍历的结束位置是原数组长度n减去当前剩余的数字，再加1，因为确保给 curNum 留够足够的位置来遍历。之后就是跳过重复，对于重复的数字，只检查一遍就好了。调用递归函数，此时的 curSum 要减去当前数字 A[i]，curNum 要减1，start 为 i+1，若递归函数返回 true，则整个返回 true。for 循环退出后返回 false。令博主感到惊讶的是，这个代码的运行速度比之前的 DP 解法还要快，叼，参见代码如下：

解法三：

class Solution {
public:
    bool splitArraySameAverage(vector<int>& A) {
        int n = A.size(), m = n / 2, sum = accumulate(A.begin(), A.end(), 0);
        bool possible = false;
        for (int i = 1; i <= m && !possible; ++i) {
            if (sum * i % n == 0) possible = true;
        }
        if (!possible) return false;
        sort(A.begin(), A.end());
        for (int i = 1; i <= m; ++i) {
            if (sum * i % n == 0 && helper(A, sum * i / n, i, 0)) return true;
        }
        return false;
    }
    bool helper(vector<int>& A, int curSum, int curNum, int start) {
        if (curNum == 0) return curSum == 0;
        if (A[start] > curSum / curNum) return false;
        for (int i = start; i < A.size() - curNum + 1; ++i) {
            if (i > start && A[i] == A[i - 1]) continue;
            if(helper(A, curSum - A[i], curNum - 1, i + 1)) return true;
        }
        return false;
    }
};