题目地址:
https://leetcode.com/problems/split-array-with-same-average/description/
给定一个长 n n n数组 A A A, A [ i ] ≥ 0 A[i]\ge 0 A[i]≥0。问是否能将其划分为两个非空子集,使得两个子集的平均数相等。 n ≤ 30 , A [ i ] ≤ 1 0 4 n\le 30, A[i]\le 10^4 n≤30,A[i]≤104。
设划分的两个子集分别为
B
,
C
B,C
B,C,则有
∑
B
∣
B
∣
=
∑
C
∣
C
∣
=
∑
B
+
∑
C
∣
B
∣
+
∣
C
∣
=
∑
A
∣
A
∣
\frac{\sum B}{|B|}=\frac{\sum C}{|C|}=\frac{\sum B+\sum C}{|B|+|C|}=\frac{\sum A}{|A|}
∣B∣∑B=∣C∣∑C=∣B∣+∣C∣∑B+∑C=∣A∣∑A而
∑
A
∣
A
∣
\frac{\sum A}{|A|}
∣A∣∑A与划分无关,所以我们其实就是要求是否存在一个子集,其平均值等于
∑
A
∣
A
∣
\frac{\sum A}{|A|}
∣A∣∑A。考虑新数组
B
[
i
]
=
A
[
i
]
−
∑
A
∣
A
∣
B[i]=A[i]-\frac{\sum A}{|A|}
B[i]=A[i]−∣A∣∑A,那么问题即转化为问
B
B
B数组是否存在和为
0
0
0的非平凡子集(非平凡的意思是不能为空,也不能为
B
B
B自己)。为了避免处理浮点数,我们可以将
B
B
B乘以一个系数,使其每个数都为整数。
这个是背包问题,通常来讲,这种题目可以用动态规划做,但此题的数据范围较为特殊,
n
n
n较小,而
A
[
i
]
A[i]
A[i]的范围较大,如果用动态规划容易爆空间。可以考虑用双向DFS来做,参考https://blog.csdn.net/qq_46105170/article/details/115587834。具体思路是,我们先考虑前
n
/
2
n/2
n/2个数,进行DFS,用哈希表存一下其所有子集的和以及每个子集的元素个数(这是为了排除掉平凡子集),如果前
n
/
2
n/2
n/2个数就已经有非平凡子集和为
0
0
0了,那就说明有解,可以提前结束;否则再暴力DFS后
n
/
2
n/2
n/2个数的每个子集的和
s
s
s,看哈希表是否存在
−
s
-s
−s并且元素个数总和大于
0
0
0小于
n
n
n,如果是,则有解。暴力枚举完毕没发现解,则无解。代码如下:
class Solution {
public:
bool splitArraySameAverage(vector<int> &A) {
int sum = 0;
for (int x : A) sum += x;
int g = gcd(sum, A.size());
int a = sum / g, b = A.size() / g;
for (int &x : A) x = x * b - a;
unordered_map<int, int> mp;
return dfs1(0, 0, 0, A, mp) || dfs2(A.size() / 2, 0, 0, A, mp);
}
bool dfs2(int u, int cnt, int sum, vector<int> &A,
unordered_map<int, int> &mp) {
auto it = mp.find(-sum);
if (it != mp.end() && it->second + cnt && it->second + cnt < A.size())
return true;
for (int i = u; i < A.size(); i++)
if (dfs2(i + 1, cnt + 1, sum + A[i], A, mp)) return true;
return false;
}
// 统计前一半的数的每个子集的和以及元素个数
bool dfs1(int u, int cnt, int sum, vector<int> &A,
unordered_map<int, int> &mp) {
if (!sum && cnt && cnt < A.size()) return true;
auto it = mp.find(sum);
if (it == mp.end() || it->second > cnt) mp[sum] = cnt;
for (int i = u; i < A.size() / 2; i++)
if (dfs1(i + 1, cnt + 1, sum + A[i], A, mp)) return true;
return false;
}
int gcd(int a, int b) { return b ? gcd(b, a % b) : a; }
};
时间复杂度 O ( 2 n / 2 + 1 ) O(2^{n/2+1}) O(2n/2+1),空间 O ( 2 n / 2 ) O(2^{n/2}) O(2n/2)。