You are given an integer array nums
.
You should move each element of nums
into one of the two arrays A
and B
such that A
and B
are non-empty, and average(A) == average(B)
.
Return true
if it is possible to achieve that and false
otherwise.
Note that for an array arr
, average(arr)
is the sum of all the elements of arr
over the length of arr
.
Example 1:
Input: nums = [1,2,3,4,5,6,7,8]
Output: true
Explanation: We can split the array into [1,4,5,8] and [2,3,6,7], and both of them have an average of 4.5.
Example 2:
Input: nums = [3,1]
Output: false
Constraints:
1 <= nums.length <= 30
0 <= nums[i] <= 104
题目链接:https://leetcode.com/problems/split-array-with-same-average/
题目大意:问能否把数组分成两部分使得两部分的平均值相等
题目分析:设分成两部分的个数和总和分别为n1, sum1, n2, sum2,则有sum1/n1 = sum2/n2等价于sum1/n1 = (sum - sum1)/(n - n1),化简得sum1/n1 = sum/n ,于是问题转化为是否能在n个数里找到n1个数,使它们的均值等于总均值,因sum1,n1是整数,故可过滤出可能的n1,又因要分成两部分,必然有min(n1, n2) <= n/2,可借此减少枚举数量,进一步将问题转化为是否能在n个数里找n1个数,使它们的和为C(C = n1 * sum / n),对于每个枚举的n1,C可当做定值,此问题可用01背包模型做,dp[i][j][k]表示枚举到i,选了j个数是否能合成k,转移方程:
dp[i][j][k] |= dp[i - 1][j - 1][k - A[i]]
dp[i][j][k] |= dp[i - 1][j][k]
641ms,时间击败40.71%
class Solution {
public boolean splitArraySameAverage(int[] A) {
int n = A.length;
if (n == 1 || (n == 2 && (A[0] != A[1]))) {
return false;
}
int sum = 0;
for (int i = 0; i < n; i++) {
sum += A[i];
}
int[] candidates = new int[n >> 1];
int cnt = 0;
for (int i = 1; i <= n / 2; i++) {
if ((i * sum) % n == 0) {
candidates[cnt++] = i;
}
}
if (cnt == 0) {
return false;
}
boolean[][][] dp = new boolean[n][n / 2 + 1][sum];
dp[0][0][0] = true;
dp[0][1][A[0]] = true;
for (int i = 1; i < n; i++) {
for (int j = 0; j <= Math.min(i + 1, n / 2); j++) {
for (int k = 0; k < sum; k++) {
if (j > 0 && k >= A[i]) {
dp[i][j][k] |= dp[i - 1][j - 1][k - A[i]];
}
dp[i][j][k] |= dp[i - 1][j][k];
// if (dp[i][j][k]) {
// System.out.println("dp["+i+"]["+j+"]["+k+"] = " + dp[i][j][k]);
// }
}
}
}
for (int i = 0; i < cnt; i++) {
int tN = candidates[i];
int tSum = (tN * sum) / n;
for (int j = 0; j < n; j++) {
if (dp[j][tN][tSum]){
return true;
}
}
}
return false;
}
}
01背包有个空间二维到一维的优化在此也适用,因每个数字A[i]只能选一次,只需将二层循环逆序即可,这样可以保证用来更新当前i阶段dp[j][k]的值必然来自i-1阶段,这里可能不太好理解,可以反过来看为什么正序枚举不对,如若正序枚举,假设当前对A[i]进行决策(选或不选),且j枚举到2,如果选A[i] => dp[2][k] |= dp[1][k-A[i]],j继续枚举到3,继续选A[i] => dp[3][k'] |= dp[2][k'- A[i]],若k = k' - A[i],即j枚举到3时选了A[i]的状态是从j枚举到2时选了A[i]的状态推导出来的,也就是说此时A[i]被选了两次,不符合题意,这就容易理解为什么j逆序枚举是正确的了,j枚举到3时,要用j等2时的dp值来更新,当前i阶段j等于2时的dp值还未被推导,故使用的必然是i-1阶段的值
188ms,时间击败61.79%
class Solution {
public boolean splitArraySameAverage(int[] A) {
int n = A.length;
if (n == 1 || (n == 2 && (A[0] != A[1]))) {
return false;
}
int sum = 0;
for (int i = 0; i < n; i++) {
sum += A[i];
}
int[] candidates = new int[n >> 1];
int cnt = 0;
for (int i = 1; i <= n / 2; i++) {
if ((i * sum) % n == 0) {
candidates[cnt++] = i;
}
}
if (cnt == 0) {
return false;
}
boolean[][] dp = new boolean[n / 2 + 1][sum];
dp[0][0] = true;
dp[1][A[0]] = true;
for (int i = 1; i < n; i++) {
for (int j = Math.min(i + 1, n / 2); j > 0; j--) {
for (int k = A[i]; k < sum; k++) {
dp[j][k] |= dp[j - 1][k - A[i]];
}
}
}
for (int i = 0; i < cnt; i++) {
int tN = candidates[i];
int tSum = (tN * sum) / n;
if (dp[tN][tSum]){
return true;
}
}
return false;
}
}
还可以进一步优化,当前的主要耗时点在于第三层的枚举其实存在大量无用的值(dp值为false的k值),如果某个阶段的dp值为false,那么之后就不可能再变成true(这是显然的),于是可以考虑直接存能组成的数字,将dp[i]变成一个hashset,表示选i个数字能组成的集合
16ms,时间击败97.17%
class Solution {
public boolean splitArraySameAverage(int[] A) {
int n = A.length;
if (n == 1 || (n == 2 && (A[0] != A[1]))) {
return false;
}
int sum = 0;
for (int i = 0; i < n; i++) {
sum += A[i];
}
boolean ok = false;
for (int i = 1; i <= n / 2; i++) {
if ((i * sum) % n == 0) {
ok = true;
break;
}
}
if (!ok) {
return false;
}
int m = n / 2 + 1;
Set<Integer>[] dp = new HashSet[m];
for (int i = 0; i < m; i++) {
dp[i] = new HashSet<>();
}
dp[0].add(0);
for (int i = 0; i < n; i++) {
for (int j = Math.min(i + 1, n / 2); j > 0; j--) {
for (int num : dp[j - 1]) {
dp[j].add(num + A[i]);
}
}
}
for (int i = 1; i < m; i++) {
if ((i * sum) % n == 0 && dp[i].contains(i * sum / n)) {
return true;
}
}
return false;
}
}