Problem statement

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

Solution 1

First of all note, that given condition is equivalent to the following: average of elements in A is equal to average of all elements. Now, if we subtract average of A from each element, we need to find such group B, which have sum of elements equal to zero. There are 2^30 possible subsets of set A, which is quite big, but we can do the following: find all possible sums for fist half and for the second half and then check if we have sum equal to zero from two halves. Note, that it can happen that we have already 0 in one of two halves, then we return True in this case immediately. If we do not have zero, than we need to find el from first half, such that -el in second half and also el is not equal to sum of elements in first half.

Complexity

Time and space complexity is O(2^{N/2}).

Code

class Solution:
    def splitArraySameAverage(self, A):
        def sums(arr):
            if not arr: return {}
            out = {arr[0]}
            for l in arr[1:]: 
                out = {z + l for z in out} | out | {l}
            return out
        
        if len(A) == 1: return False
        
        n, S = len(A), sum(A)
        A = [z*n - S for z in A]
        part1, part2 = sums(A[:n//2]), sums(A[n//2:])
        sum1, sum2 = sum(A[:n//2]), sum(A[n//2:])
    
        if 0 in part1 or 0 in part2: return True

        for el in part1:
            if el != sum1 and -el in part2:
                return True
        
        return False

Solution 2

There is another O(n^2 * m) solution, using classical dp: in sums[i] we will keep all possible sums if we allowed to use exaclty i elements.

Complexity

It is O(n^2 * m) for time and O(n*m) space.

Code 1

class Solution:
    def splitArraySameAverage(self, A):
        sums = defaultdict(set)
        sums[0].add(0)
        n = len(A)
        
        for i in range(n):
            for j in range(i+1, -1, -1):
                for elem in sums[j]: sums[j+1].add(elem + A[i])
        
        return any(sum(A)*j % n == 0 and sum(A)*j//n in sums[j] for j in range(1, n))

Code 2

The same idea can be implemented using binary masks, see Problem 0416 Partition Equal Subset Sum with similar idea.

#### borrowed code
class Solution:
    def splitArraySameAverage(self, A):
    		N, S, P = len(A), sum(A), [1]
    		for a in A: P[1:] = [(p << a) | q for p, q in zip(P, P[1:] + [0])]
    		return any(S * n % N == 0 and P[n] & (1 << (S * n // N)) for n in range(1, N))