1.lc 805. 数组的均值分割(2022/11/14)
题目:
给定你一个整数数组 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。
思路:【官方题解】
作者:力扣官方题解
链接:https://leetcode.cn/problems/split-array-with-same-average/solutions/1966163/shu-zu-de-jun-zhi-fen-ge-by-leetcode-sol-f630/
来源:力扣(LeetCode)
设数组 nums 的长度为 n,假设我们移动了 k 个元素到数组 A 中,移动了 n−k 个元素到数组 B 中,此时两个数组的平均值相等。设 sum(A),sum(B),sum(nums) 分别表示数组 A,B,nums 的元素和,由于数组 A,B 的平均值相等,则有 ,上述等式可以进行变换为如下:
sum(A)×(n−k)=sum(B)×k
⇔ sum(A)×n=(sum(A)+sum(B))×k
⇔ sum(A)×n=(sum(nums))×k
设数组 nums 的平均数为 ,我们移动了 k 个元素到数组 A 中,则此时已经数组 A 的平均数也为 avg,则此时数组 A 的元素之和为 sum(A)=k×avg,则该问题可以等价于在数组中取 k 个数,使得其和为 k×avg。对应的「0-1背包问题」则为是否可以取 k 件物品使得背包的重量为 k×avg。我们设 dp[i][x] 表示当前已从数组 nums 取出 i 个元素构成的和为 x 的可能性:
如果 dp[i][x]=true,表示当前已经遍历过的元素中可以取出 i 个元素构成的和为 x;
如果 dp[i][x]=false,表示当前已经遍历过的元素中不存在取出 i 个元素的和等于 x;
假设前 j−1 个元素中存在长度为 i 的子集且子集的和为 x,则此时 dp[i][x]=true,我们当前遍历 nums[j] 时,则可以推出一定存在长度为 i+1 的子集且满足子集的和为 x+nums[j],可以得到状态转移方程为:dp[i+1][x+nums[j]]=dp[i][x]
根据题意可知:,可以变换为: sum(A)×n=sum(nums)×k,所以我们只需要找到一个元素个数为 k 的子集 A 满足上述条件即可,因此我们利用上述递推公式求出所有可能长度的子集和组合即可,此时需要的时间复杂度约为
,按照题目给定的测试用例可能会超时,所以需要一些减枝技巧,具体的减枝技巧如下:
根据题意可以推出 ,则此时如果满足题目要求,则一定存在整数 k∈[1,n],使得 [sum(nums)×k] mod n=0,因此我们可以提前是否存在 k 满足上述要求,如果不存在则可以提前终止。
根据题目要求可以划分为两个子数组 A,B,则可以知道两个子数组中一定存在一个数组的长度小于等于 ,因此我们只需检测数组长度小于等于
的子数组中是否存在其和满 subsum×n=nums×k 即可。
代码:
class Solution {
public:
bool splitArraySameAverage(vector<int>& nums) {
int n = nums.size(), m = n / 2, i;
int sum = accumulate(nums.begin(), nums.end(), 0);
for(i = 1; i <= m; i++){
if(sum * i / n == 0) break;
}
if(i > m) return false;
vector<unordered_set<int>>dp(m + 1);
dp[0].insert(0);
for(int num: nums){
for(int i = m; i >= 1; i--){
for(int x: dp[i - 1]){
int cur = x + num;
if(cur * n == sum * i) return true;
dp[i].insert(cur);
}
}
}
return false;
}
};