[LeetCode] 805. Split Array With Same Average 用相同均值拆分数组
In a given integer array A, we must move every element of A to either list B or list C. (B and C initially start empty.)
Return true if and only if after such a move, it is possible that the average value of B is equal to the average value of C, and B and C are both non-empty.
Example : Input: [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 the average of 4.5.
Note:
- The length of
Awill be in the range [1, 30]. A[i]will be in the range of[0, 10000].
给一个数组A，把A中的每一个元素都移到数组B或C中(B, C初始为空)。如果移动后可以使B和C的均值相等，则返回ture。其实就是将一个数组分成两部分，每部分的平均值相同。
题目的关键点是：当能拆分成两个平均值相等的数组时，拆分的数组和原来数组的平均值是相同的。因此问题转换为，先算出A的平均值A_aver，如果A中的数组成的子数组B的平均值等于A_aver，再去判断剩余的数组成的数组C的平均值是否和A_aver相等。
The key thing of this problem is, when we are able to make a same average split, the average of each splitted array is the same as the average of the whole array.
So the problem can be transformed into a simpler one: given a target number (tosum), can we construct it using a specific number (lenB) of integers in a list(A).
Then we can try every possible numbers of lenB, to see whether any one is feasible.
If the array of size n can be splitted into group A and B with same mean, assuming A is the smaller group, then
totalSum/n = Asum/k = Bsum/(n-k), where k = A.size() and 1 <= k <= n/2;
Asum = totalSum*k/n, which is an integer. So we have totalSum*k%n == 0;
如果一个长度为n的数组可以被划分为A和B两个数组，假设A的长度小于B并且A的大小是k，那么：total_sum / n == A_sum / k == B_sum / (n - k)，其中1 <= k <= n / 2。可得出：A_sum = total_sum * k / n。由于A_sum一定是个整数，所以可以推导出total_sum * k % n == 0，那就是说，对于特定的total_sum和n而言，符合条件的k不会太多。首先验证是否存在符合条件的k，如果不存在就可以提前返回false。
解法1: early pruning + knapsack DP, O(n^3 * M)
如果经过第一步的验证，发现确实有符合条件的k，那么我们在第二步中，就试图产生k个子元素的所有组合，并且计算他们的和。这里的思路就有点类似于背包问题了，我们的做法是：定义vector<vector<unordered_set<int>>> sums，其中sums[i][j]表示A[0, i]这个子数组中的任意j个元素的所有可能和。可以得到递推公式是：sums[i][j] = sums[i - 1][j] "join" (sums[i][j - 1] + A[i])，其中等式右边的第一项表示这j个元素中不包含A[i]，而第二项表示这j个元素包含A[i]。这样就可以采用动态规划的思路得到sums[n - 1][k]了（1 <= k <= n / 2）。
有了sums[n - 1][k]，我们就检查sums[n - 1][k]中是否包含(total_sum * k / n)。一旦发现符合条件的k，就返回true，否则就返回false。
If there are still some k valid after early pruning by checking totalSum*k%n == 0,
we can generate all possible combination sum of k numbers from the array using DP, like knapsack problem. (Note: 1 <= k <= n/2)
Next, for each valid k, simply check whether the group sum, i.e. totalSum * k / n, exists in the kth combination sum hashset.
vector<vector<unordered_set<int>>> sums(n, vector<unordered_set<int>>(n/2+1));
sums[i][j] is all possible combination sum of j numbers from the subarray A[0, i];
Goal: sums[n-1][k], for all k in range [1, n/2]
Initial condition: sums[i][0] = {0}, 0 <= i <= n-1; sums[0][1] = {all numbers in the array};
Deduction: sums[i+1][j] = sums[i][j] "join" (sums[i][j-1] + A[i+1])
The following code uses less space but the same DP formula.
Runtime analysis:
All numbers in the array are in range [0, 10000]. Let M = 10000.
So the size of kth combination sum hashset, i.e. sums[...][k], is <= k * M;
For each number in the array, the code need loop through all combination sum hashsets, so
the total runtime is n * (1 * M + 2 * M + ... + (n/2) * M) = O(n^3 * M)
解法2: TLE, For such k, the problem transforms to "Find k sum = Asum, i.e. totalSum * k/n, from an array of size n". This subproblem is similar to LC39 combination sum, which can be solved by backtracking.
Python:
class Solution(object):
def splitArraySameAverage(self, A):
if len(A)==1: return False
global_avg = sum(A)/float(len(A))
for lenB in range(1, len(A)/2+1):
if int(lenB*global_avg) == lenB*global_avg:
if self.exist(lenB*global_avg, lenB, A):
return True
return False
def exist(self, tosum, item_count, arr):
if item_count==0:
return False if tosum else True
if item_count > len(arr) or not arr:
return False
if any([self.exist(tosum-arr[0], item_count-1, arr[1:]),
self.exist(tosum, item_count, arr[1:])]):
return True
return False
Python:
# Time: O(n^4)
# Space: O(n^3)
class Solution(object):
def splitArraySameAverage(self, A):
"""
:type A: List[int]
:rtype: bool
"""
def possible(total, n):
for i in xrange(1, n//2+1):
if total*i%n == 0:
return True
return False
n, s = len(A), sum(A)
if not possible(n, s):
return False
sums = [set() for _ in xrange(n//2+1)];
sums[0].add(0)
for num in A: # O(n) times
for i in reversed(xrange(1, n//2+1)): # O(n) times
for prev in sums[i-1]: # O(1) + O(2) + ... O(n/2) = O(n^2) times
sums[i].add(prev+num)
for i in xrange(1, n//2+1):
if s*i%n == 0 and s*i//n in sums[i]:
return True
return False
C++: 1
class Solution {
public:
bool splitArraySameAverage(vector<int>& A) {
int n = A.size(), m = n/2, totalSum = accumulate(A.begin(), A.end(), 0);
// early pruning
bool isPossible = false;
for (int i = 1; i <= m && !isPossible; ++i)
if (totalSum*i%n == 0) isPossible = true;
if (!isPossible) return false;
// DP like knapsack
vector<unordered_set<int>> sums(m+1);
sums[0].insert(0);
for (int num: A) {
for (int i = m; i >= 1; --i)
for (const int t: sums[i-1])
sums[i].insert(t + num);
}
for (int i = 1; i <= m; ++i)
if (totalSum*i%n == 0 && sums[i].find(totalSum*i/n) != sums[i].end()) return true;
return false;
}
};
C++: 1
class Solution {
public:
bool splitArraySameAverage(vector<int>& A) {
int n = A.size(), m = n / 2;
int totalSum = accumulate(A.begin(), A.end(), 0);
// early pruning
bool isPossible = false;
for (int i = 1; i <= m; ++i) {
if (totalSum * i % n == 0) {
isPossible = true;
break;
}
}
if (!isPossible) {
return false;
}
// DP like knapsack
vector<unordered_set<int>> sums(m + 1);
sums[0].insert(0);
for (int num: A) { // for each element in A, we try to add it to sums[i] by joining sums[i - 1]
for (int i = m; i >= 1; --i) {
for (const int t: sums[i - 1]) {
sums[i].insert(t + num);
}
}
}
for (int i = 1; i <= m; ++i) {
if (totalSum * i % n == 0 && sums[i].find(totalSum * i / n) != sums[i].end()) {
return true;
}
}
return false;
}
};
C++: 2 TLE
class Solution {
public:
bool splitArraySameAverage(vector<int>& A) {
int n = A.size(), m = n/2, totalSum = accumulate(A.begin(), A.end(), 0);
sort(A.rbegin(), A.rend()); // Optimization
for (int i = 1; i <= m; ++i)
if (totalSum*i%n == 0 && combinationSum(A, 0, i, totalSum*i/n)) return true;
return false;
}
bool combinationSum(vector<int>& nums, int idx, int k, int tar) {
if (tar > k * nums[idx]) return false; // Optimization, A is sorted from large to small
if (k == 0) return tar == 0;
for (int i = idx; i <= nums.size()-k; ++i)
if (nums[i] <= tar && combinationSum(nums, i+1, k-1, tar-nums[i])) return true;
return false;
}
};
