第六章 - 高频考题(困难)
0805. 数组的均值分割

题目地址(805. 数组的均值分割)

https://leetcode-cn.com/problems/split-array-with-same-average/

题目描述

1
给定的整数数组 A ,我们要将 A数组 中的每个元素移动到 B数组 或者 C数组中。(B数组和C数组在开始的时候都为空)
2
3
返回true ,当且仅当在我们的完成这样的移动后,可使得B数组的平均值和C数组的平均值相等,并且B数组和C数组都不为空。
4
5
示例:
6
输入:
7
[1,2,3,4,5,6,7,8]
8
输出: true
9
解释: 我们可以将数组分割为 [1,4,5,8] 和 [2,3,6,7], 他们的平均值都是4.5。
10
11
12
注意:
13
14
A 数组的长度范围为 [1, 30].
15
A[i] 的数据范围为 [0, 10000].
Copied!

前置知识

公司

  • 暂无

思路

实际上分出的两个列表 B 和 C 的均值都等于列表 A 的均值,这是本题的入手点。以下是证明:
令 B 的长度为 K,A 的长度为 N。 则有 sum(B)/K = sum(C)/(N-K)。
进而:
1
sum(B) * (N - K) = sum(C) * K
2
sum(B) * N = (sum(B) + sum(C)) * K
3
sum(B) / K = (sum(B) + sum(C)) / N
4
sum(B) / K = sum(A) / N
Copied!
因此我们可以枚举所有的 A 的大小 i,相应地 B 的大小就是 n - i,其中 n 为数组 A 的大小。
而由于两个列表 B 和 C 的均值都等于列表 A 的均值。因此可以提前计算出 A 的均值 avg,那么 A 的总和其实就是 i * avg ,我们使用回溯找到一个和为 i * avg 的组合,即可返回 true,否则返回 false。
值得注意的是,我们只需要枚举 i 为 1 到 N//2 范围即可,这可以达到剪枝的效果。
核心代码:
1
def splitArraySameAverage(self, A: List[int]) -> bool:
2
n = len(A)
3
avg = sum(A) / n
4
5
for i in range(1, n // 2 + 1):
6
for combination in combinations(A, i):
7
if abs(sum(combination) - avg * i) < 1e-6:
8
return True
9
return False
Copied!
上面代码由于回溯里面嵌套了 sum,因此时间复杂度为回溯的时间复杂度 * sum 的时间复杂度,因此总的时间复杂度在最坏的情况下是 $n * 2^n$。代入题目的 n 范围是 30,一般这种复杂度只能解决 20 以下的题目,因此需要考虑优化。
我们可以不计算出来所有的组合之后再求和,而是直接计算所有的和的组合,这种算法的时间复杂度为 $2^n$。
核心代码:
1
def splitArraySameAverage(self, A: List[int]) -> bool:
2
n = len(A)
3
avg = sum(A) / n
4
5
for i in range(1, n // 2 + 1):
6
for s in combinationSum(A, i):
7
if abs(s - avg * i) < 1e-6:
8
return True
9
return False
Copied!
但是遗憾的是,这仍然不足以通过所有的测试用例。
接下来,我们可以通过进一步剪枝的手段来达到 AC 的目的。 很多回溯的题目都是基于剪枝来完成的。剪枝是回溯问题的核心考点。
这个技巧就是双向搜索,双向搜索相比之前的回溯可达到减少指数数字的效果,从 $O(2^n)$ 降低到 $O(2^(N//2))$。代入题目,这样指数变为了 30/2 = 15,就可以通过了。
具体地,我们可以 combinationSum A 数组的一半(不妨称 A1),然后 combinationSum A 数组的令一半(不妨称 A2),那么 A1 和 A2 的总和如果是 avg * i 不也行么?简单起见,我们可以令 A1 为数组 A 的前一半, A2 为数组的后一半。
同时,为了避免这种加法,我们可以对问题进行一个转化。即将数组 A 的所有数都减去 avg,这样问题转化为找到一个和为 0 的组合,即可以找到一个和为 avg * i 的组合。

关键点

  • 双端搜索

代码

  • 语言支持:Python3
Python3 Code:
1
class Solution(object):
2
def splitArraySameAverage(self, A):
3
from fractions import Fraction
4
N = len(A)
5
total = sum(A)
6
A = [a - Fraction(total, N) for a in A] # 转化后的 A,免于计算 sum
7
8
if N == 1: return False
9
10
S1 = {A[0]} # 所有 A1 可能的和的集合
11
for i in range(1, N//2):
12
# {a + A[i] for a in S1} 在之前选择的基础上选择 A[i] 的新集合
13
# {A[i]} 是仅选择 A[i] 的新集合
14
# S1 是不选择 A[i] 的集合
15
# | 是集合并操作
16
S1 = {a + A[i] for a in S1} | S1 | {A[i]}
17
if 0 in S1: return True
18
19
S2 = {A[-1]} # 所有 A2 可能的和的集合
20
for i in range(N//2, N-1):
21
S2 = {a + A[i] for a in S2} | S2 | {A[i]}
22
if 0 in S2: return True
23
# 为了避免 B 选择了所有的 A,这样 C 就是空了,我们增加一个这样的判断
24
sleft = sum(A[i] for i in range(N//2))
25
sright = sum(A[i] for i in range(N//2, N))
26
27
return any(-ha in S2 and (ha, -ha) != (sleft, sright) for ha in S1)
28
Copied!
复杂度分析
令 n 为数组长度。
  • 时间复杂度:$O(2^(N//2))$
  • 空间复杂度:$O(2^(N//2))$
此题解由 力扣刷题插件 自动生成。
力扣的小伙伴可以关注我,这样就会第一时间收到我的动态啦~
以上就是本文的全部内容了。大家对此有何看法,欢迎给我留言,我有时间都会一一查看回答。更多算法套路可以访问我的 LeetCode 题解仓库:https://github.com/azl397985856/leetcode 。 目前已经 40K star 啦。大家也可以关注我的公众号《力扣加加》带你啃下算法这块硬骨头。
关注公众号力扣加加,努力用清晰直白的语言还原解题思路,并且有大量图解,手把手教你识别套路,高效刷题。