折半搜索
805. 数组的均值分割
给定你一个整数数组 nums
我们要将 nums 数组中的每个元素移动到 A 数组 或者 B 数组中，使得 A 数组和 B 数组不为空，并且 average(A) == average(B) 。
如果可以完成则返回true ， 否则返回 false 。
注意：对于数组 arr , average(arr) 是 arr 的所有元素除以 arr 长度的和。
数据范围：\(1 \le nums.length \le 30\), \(0 \le nums[i] \le 10^4\)
方法一 折半搜索
首先，如果 \(k\) 个元素分给 \(A\)，剩余 \(n - k\) 个分给 \(B\)，设 \(A\) 的平均值是 \(avg_A\)，\(B\) 的平均值是 \(avg_B\)，那么有：
所以，\(A\) 和 \(B\) 的平均值即整体的平均值。如果将所有元素变为 \(nums[i] - avg\)，那么问题等价于分成 \(A\) 和 \(B\) 两个非空数组，每个数组的和是 \(0\)。
这个问题，如果直接每个元素的分配情况，那么复杂度是 \(O(2^n)\)，而 \(n\) 最大到 \(30\)，所以无法通过本题。
我们采用折半枚举的方式，首先枚举前一半数字的分配方式，将所有可以分配到 \(A\) 中的和放到一个集合 \(st\) 中。然后再枚举后一半数字的分配方式，如果凑出一中方案使得分配到 \(A\) 的中是 \(x\)，而 \(-x\) 刚好在 \(st\) 中，那么答案就是 \(true\)。
最后，需要注意 \(avg\) 可能是小数，这给我们处理带来了麻烦，但是我们可以提前给所有数字乘一个 \(n\)，这样就可以保证不影响问题求解的前提下，使得 \(avg\) 一定是整数。
注意特殊情况，\(n = 1\)，以及 \(A\) 和 \(B\) 非空。
class Solution {
public:
bool splitArraySameAverage(vector<int>& nums) {
int n = nums.size(), sum = accumulate(nums.begin(), nums.end(), 0) * n;
if (n == 1) {
return false;
}
unordered_set<int> st;
int m = n / 2, avg = sum / n;
// 为了保证一定有数字被选中，i 从 1 开始
for (int i = 1; i < (1 << m); i++) {
int cur = 0;
for (int j = 0; j < m; j++) {
if (i >> j & 1) {
cur += nums[j] * n - avg;
}
}
if (cur == 0) {
return true;
}
st.insert(cur);
}
// 为了保证一定有数字不被选中，i 不能到最大值
for (int i = 0; i < (1 << (n - m)) - 1; i++) {
int cur = 0;
for (int j = 0; j < (n - m); j++) {
if (i >> j & 1) {
cur += nums[m + j] * n - avg;
}
}
if (st.find(-cur) != st.end()) {
return true;
}
}
return false;
}
};
方法二 动态规划
类似背包的思想，选择 \(i\) 个物品，总和是否到达 \(j\)。那么转移时，对于数字 \(x\)：
初始时，\(d[0][0] = 1\)。过程中，如果 \(d[i][j] = 1\)，并且 \(j / i = sum / n\)，即 \(j * n = sum * i\)，那么答案为 \(true\)。
但是这样做负载度为 \(O(n^2sum)\)，最高时为 \(6\times 10^9\),在力扣上无法安全通过，因此需要加一些剪枝优化：
- 过程中，有很多状态是无效的，因此 \(d[i]\) 表示为一个集合，直接对其中的元素进行转移。
- 因为总有一个集合包含的元素个数小于 \(n / 2\)，因此，最多只需考虑转移到 \(d[n / 2]\) 即可。
- 预处理，考虑分给 \(A\) 数组 \(i\) 个元素，如果所有的 \(i\) 都不满足 \(sum * i \bmod n\) 是 \(0\) ，那么永远无法选出一组数字，他们的平均数是 \(sum / n\)。
class Solution {
public:
bool splitArraySameAverage(vector<int>& nums) {
int sum = accumulate(nums.begin(), nums.end(), 0);
int n = nums.size();
if (n == 1) {
return false;
}
int m = n / 2;
bool fg = false;
for (int i = 1; i <= m; i++) {
if (sum * i % n == 0) {
fg = true;
break;
}
}
if (!fg) return false;
vector<unordered_set<int>> d(m + 1);
d[0].insert(0);
for (int i = 1; i <= n; i++) {
int x = nums[i - 1];
for (int j = m; j >= 1; j--) {
for (auto k : d[j - 1]) {
k += x;
d[j].insert(k);
if (sum * j == k * n) {
return true;
}
}
}
}
return false;
}
};