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; } };
Github 同步地址:
https://github.com/grandyang/leetcode/issues/805
类似题目:
参考资料:
https://leetcode.com/problems/split-array-with-same-average/