LeetCode 805: Split Array With Same Average

Problem Restatement

We are given an integer array nums.

We need to split every element into one of two arrays, A and B, such that:

A is non-empty
B is non-empty
average(A) == average(B)

Return True if such a split is possible. Otherwise, return False.

The array length is at most 30, and each value is between 0 and 10000. The official examples include nums = [1,2,3,4,5,6,7,8], which returns true, and nums = [3,1], which returns false.

Input and Output

Examples

Example 1:

nums = [1, 2, 3, 4, 5, 6, 7, 8]

One valid split is:

A = [1, 4, 5, 8]
B = [2, 3, 6, 7]

Both arrays have sum 18 and length 4, so both averages are 4.5.

The answer is:

True

Example 2:

nums = [3, 1]

The only possible split is:

A = [3]
B = [1]

Their averages are different.

The answer is:

False

First Thought: Brute Force

A direct approach is to try every subset as A.

The remaining elements become B.

For each subset, check whether:

sum(A) / len(A) == sum(B) / len(B)

This works, but there are 2^n subsets.

Since n can be 30, this is too large.

Key Insight

We do not need to build both arrays.

Let:

n = len(nums)
total = sum(nums)

Suppose A has k elements and sum s.

For A and B to have the same average, A must also have the same average as the whole array:

s / k == total / n

So:

s = k * total / n

This means the problem becomes:

Find a non-empty subset of size k whose sum is exactly k * total / n.

We only need to check subset sizes from 1 to n // 2.

If a valid larger subset exists, its complement is a valid smaller subset.

Algorithm

Create a list of sets:

dp[k] = all sums possible using exactly k elements

Initialize:

dp[0] = {0}

Then process every number in nums.

For each number num, update subset sizes in reverse:

for k in range(n // 2, 0, -1):
    for prev_sum in dp[k - 1]:
        dp[k].add(prev_sum + num)

Reverse order matters because each number can be used only once.

After building possible sums, check every subset size k from 1 to n // 2.

For a valid target sum, this must be divisible:

k * total % n == 0

If divisible, compute:

target = k * total // n

If target is in dp[k], return True.

If no subset size works, return False.

Correctness

The dynamic programming state dp[k] contains exactly the sums that can be formed by choosing k elements from the numbers already processed.

This is true initially because dp[0] = {0} represents choosing no elements with sum zero.

When processing a new number num, every old sum in dp[k - 1] can form a new sum in dp[k] by adding num. Updating k in reverse order ensures the same num is not reused in the same iteration.

So after all numbers are processed, dp[k] contains exactly all sums of all subsets of size k.

A valid split exists exactly when there is a subset A with size k and sum s such that:

s / k == total / n

This is equivalent to:

s = k * total / n

The algorithm checks every possible smaller-side size k and tests whether this exact target sum is reachable.

If the algorithm returns True, such a subset exists, and its complement gives the second non-empty array with the same average.

If the algorithm returns False, no subset of any valid size has the required sum, so no valid split exists.

Complexity

Let S = sum(nums).

In practice, using sets is usually compact enough for the given constraints.

Implementation

class Solution:
    def splitArraySameAverage(self, nums: list[int]) -> bool:
        n = len(nums)
        total = sum(nums)

        possible_size_exists = False
        for k in range(1, n // 2 + 1):
            if (k * total) % n == 0:
                possible_size_exists = True
                break

        if not possible_size_exists:
            return False

        dp = [set() for _ in range(n // 2 + 1)]
        dp[0].add(0)

        for num in nums:
            for k in range(n // 2, 0, -1):
                for prev_sum in list(dp[k - 1]):
                    dp[k].add(prev_sum + num)

        for k in range(1, n // 2 + 1):
            if (k * total) % n != 0:
                continue

            target = (k * total) // n

            if target in dp[k]:
                return True

        return False

Code Explanation

We first compute:

n = len(nums)
total = sum(nums)

Before building DP, we do a quick divisibility check:

if (k * total) % n == 0:

If no subset size can produce an integer target sum, then a valid split is impossible.

The DP array stores possible sums by subset size:

dp = [set() for _ in range(n // 2 + 1)]
dp[0].add(0)

For each number, we update from larger subset sizes down to smaller ones:

for k in range(n // 2, 0, -1):

This prevents using the same number more than once.

For every sum that was possible with k - 1 elements, we can create a new sum with k elements:

dp[k].add(prev_sum + num)

After processing all numbers, we test each subset size.

If the required sum exists in dp[k], then a valid subset exists:

if target in dp[k]:
    return True

Otherwise, no valid split exists:

return False

Testing

def run_tests():
    s = Solution()

    assert s.splitArraySameAverage([1, 2, 3, 4, 5, 6, 7, 8]) is True
    assert s.splitArraySameAverage([3, 1]) is False
    assert s.splitArraySameAverage([1, 2, 3]) is True
    assert s.splitArraySameAverage([5, 5, 5, 5]) is True
    assert s.splitArraySameAverage([0]) is False
    assert s.splitArraySameAverage([0, 0]) is True

    print("all tests passed")

run_tests()