题目描述
给定你一个整数数组 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 <= 300 <= nums[i] <= 104
解法
方法一：折半查找 + 二进制枚举
根据题目要求，要判断是否可以将数组 $\textit{nums}$ 划分为两个子数组 $A$ 和 $B$，使得两个子数组的平均值相等。
我们记数组 $\textit{nums}$ 的和为 $s$，元素个数为 $n$。子数组 $A$ 的和以及个数分别为 $s_1$ 和 $k$，那么子数组 $B$ 的和为 $s_2 = s - s_1$，个数为 $n - k$，即：
$$
\frac{s_1}{k} = \frac{s_2}{n - k} = \frac{s-s_1}{n-k}
$$
整理可得：
$$
s_1 \times (n-k) = (s-s_1) \times k
$$
化简可得：
$$
\frac{s_1}{k} = \frac{s}{n}
$$
也就是说，要我们找出一个子数组 $A$，使得其平均值等于数组 $\textit{nums}$ 的平均值。我们考虑将数组 $\textit{nums}$ 每个元素都减去数组 $\textit{nums}$ 的平均值，这样问题就转化为了在数组 $\textit{nums}$ 中找出一个子数组，使得其和为 $0$。
但是，数组 $\textit{nums}$ 的平均值可能不是整数，浮点数计算可能存在精度问题，我们不妨将数组 $\textit{nums}$ 中的每个元素都乘以 $n$，即 $nums[i] \leftarrow nums[i] \times n$，上述式子就变成：
$$
\frac{s_1\times n}{k} = s
$$
此时我们将数组 $\textit{nums}$ 中每个元素都减去整数 $s$，题目就转化为：在数组 $nums$ 中找出一个子数组 $A$，使得其和为 $0$。
数组 $\textit{nums}$ 的长度范围为 $[1, 30]$，如果我们使用暴力枚举子数组的方法，时间复杂度为 $O(2^n)$，会超时。我们可以使用折半查找的方法，将时间复杂度降低到 $O(2^{n/2})$。
我们将数组 $\textit{nums}$ 分成左右两部分，那么子数组 $A$ 可能存在三种情况：
- 子数组 $A$ 完全在数组 $\textit{nums}$ 的左半部分；
- 子数组 $A$ 完全在数组 $\textit{nums}$ 的右半部分；
- 子数组 $A$ 一部分在数组 $\textit{nums}$ 的左半部分，一部分在数组 $\textit{nums}$ 的右半部分。
我们可以使用二进制枚举的方法，先枚举左半部分所有子数组的和，如果存在一个子数组和为 $0$，直接返回 true，否则我们将其存入哈希表 $\textit{vis}$ 中；然后枚举右半部分所有子数组的和，如果存在一个子数组和为 $0$，直接返回 true，否则我们判断此时哈希表 $\textit{vis}$ 中是否存在该和的相反数，如果存在，直接返回 true。
需要注意的是，我们不能同时全选左半部分和右半部分，因为这样会导致子数组 $B$ 为空，这是不符合题意的。在实现上，我们只需要考虑数组的 $n-1$ 个数。
时间复杂度 $O(n\times 2^{\frac{n}{2}})$，空间复杂度 $O(2^{\frac{n}{2}})$。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20 | class Solution:
def splitArraySameAverage(self, nums: List[int]) -> bool:
n = len(nums)
if n == 1:
return False
s = sum(nums)
for i, v in enumerate(nums):
nums[i] = v * n - s
m = n >> 1
vis = set()
for i in range(1, 1 << m):
t = sum(v for j, v in enumerate(nums[:m]) if i >> j & 1)
if t == 0:
return True
vis.add(t)
for i in range(1, 1 << (n - m)):
t = sum(v for j, v in enumerate(nums[m:]) if i >> j & 1)
if t == 0 or (i != (1 << (n - m)) - 1 and -t in vis):
return True
return False
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38 | class Solution {
public boolean splitArraySameAverage(int[] nums) {
int n = nums.length;
if (n == 1) {
return false;
}
int s = Arrays.stream(nums).sum();
for (int i = 0; i < n; ++i) {
nums[i] = nums[i] * n - s;
}
int m = n >> 1;
Set<Integer> vis = new HashSet<>();
for (int i = 1; i < 1 << m; ++i) {
int t = 0;
for (int j = 0; j < m; ++j) {
if (((i >> j) & 1) == 1) {
t += nums[j];
}
}
if (t == 0) {
return true;
}
vis.add(t);
}
for (int i = 1; i < 1 << (n - m); ++i) {
int t = 0;
for (int j = 0; j < (n - m); ++j) {
if (((i >> j) & 1) == 1) {
t += nums[m + j];
}
}
if (t == 0 || (i != (1 << (n - m)) - 1) && vis.contains(-t)) {
return true;
}
}
return false;
}
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25 | class Solution {
public:
bool splitArraySameAverage(vector<int>& nums) {
int n = nums.size();
if (n == 1) return false;
int s = accumulate(nums.begin(), nums.end(), 0);
for (int& v : nums) v = v * n - s;
int m = n >> 1;
unordered_set<int> vis;
for (int i = 1; i < 1 << m; ++i) {
int t = 0;
for (int j = 0; j < m; ++j)
if (i >> j & 1) t += nums[j];
if (t == 0) return true;
vis.insert(t);
}
for (int i = 1; i < 1 << (n - m); ++i) {
int t = 0;
for (int j = 0; j < (n - m); ++j)
if (i >> j & 1) t += nums[m + j];
if (t == 0 || (i != (1 << (n - m)) - 1 && vis.count(-t))) return true;
}
return false;
}
};
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39 | func splitArraySameAverage(nums []int) bool {
n := len(nums)
if n == 1 {
return false
}
s := 0
for _, v := range nums {
s += v
}
for i, v := range nums {
nums[i] = v*n - s
}
m := n >> 1
vis := map[int]bool{}
for i := 1; i < 1<<m; i++ {
t := 0
for j, v := range nums[:m] {
if (i >> j & 1) == 1 {
t += v
}
}
if t == 0 {
return true
}
vis[t] = true
}
for i := 1; i < 1<<(n-m); i++ {
t := 0
for j, v := range nums[m:] {
if (i >> j & 1) == 1 {
t += v
}
}
if t == 0 || (i != (1<<(n-m))-1 && vis[-t]) {
return true
}
}
return false
}
|