[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 A will 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 LC39combination 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;
    }
};

　　

All LeetCode Questions List 题目汇总