Split Array With Same Average

How it works

We can build all achievable sums for each subset size using dynamic programming. For each element, we update what sums are achievable for each size. We iterate backwards through sizes to avoid using the same element twice. Finally, we check if any valid target sum exists for sizes 1 through n/2.

Steps

1. Create dp[a] as a set of achievable sums for subsets of size a. Initialize dp[0] = {0}.
2. For each number in the array:
   • For sizes from n/2 down to 1:
     • For each previously achievable sum in dp[a-1], add sum + num to dp[a].
3. After processing all numbers, check each size a from 1 to n/2:
   • If a * total % n == 0 and the target sum a * total / n exists in dp[a], return true.
4. Return false if no valid partition exists.

Complexity

• Time complexity: O(n ^ 2 * s)
• Space complexity: O(n ^ 2 * s)

Where n is the size of the input array nums, and s is the sum of the elements of the array.

Things to watch for

Forgetting Both Groups Must Be Non-Empty

A common mistake is checking only whether the averages match without ensuring both groups contain at least one element. If all elements go into group A and group B is empty, that is not a valid split. Always verify that both partitions have size greater than zero before comparing averages.

Using Floating-Point Comparison for Averages

Comparing averages directly using floating-point division can lead to precision errors. Instead of checking sum(A) / len(A) == sum(B) / len(B), use the mathematically equivalent integer comparison sum(A) * len(B) == sum(B) * len(A) to avoid floating-point inaccuracies.

Missing the Early Termination Condition

The condition a * total % n == 0 is crucial for pruning. For a valid subset of size a, the required sum a * total / n must be an integer. Skipping this check leads to unnecessary computation for impossible subset sizes and can cause the solution to time out.

Integer Overflow in Sum Calculations

When computing sums or products like a * total, the result can exceed 32-bit integer limits for large arrays with large values. Use 64-bit integers (long/long long) for intermediate calculations to prevent overflow and incorrect results.

Iterating Over All Subset Sizes Instead of Half

Since finding a valid subset A automatically determines a valid subset B (the remaining elements), you only need to check subset sizes from 1 to n/2. Checking sizes beyond n/2 is redundant because if a subset of size a works, so does the complement of size n - a.

from typing import List

class Solution:
    def splitArraySameAverage(self, nums: List[int]) -> bool:
        n = len(nums)
        # len(A) = a, len(B) = b, let a <= b
        # avg(A) = avg(B)
        # sum(A) / a = sum(B) / b = sum(nums) / n
        # sum(A) / a = avg => sum(A) = a * avg
        # sum(A) = a * sum(nums) / n
        # Find if any subset exists with a * sum(nums) / n
        # a is in the range [1, (n//2)]

        total = sum(nums)
        dp = [set() for _ in range(n // 2 + 1)]

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

        for a in range(1, n // 2 + 1):
            if (a * total % n == 0) and (a * total // n) in dp[a]:
                return True

        return False

class Solution {
public:
    bool splitArraySameAverage(vector<int>& nums) {
        int n = nums.size();

        // len(A) = a, len(B) = b, let a <= b
        // avg(A) = avg(B)
        // sum(A) / a = sum(B) / b = sum(nums) / n
        // sum(A) / a = avg => sum(A) = a * avg
        // sum(A) = a * sum(nums) / n
        // Find if any subset exists with a * sum(nums) / n
        // a is in the range [1, (n//2)]

        int total = accumulate(nums.begin(), nums.end(), 0);
        vector<unordered_set<int>> dp(n / 2 + 1);
        dp[0].insert(0);

        for (int num : nums) {
            for (int a = n / 2; a >= 1; a--) {
                for (int prev : dp[a - 1]) {
                    dp[a].insert(prev + num);
                }
            }
        }

        for (int a = 1; a <= n / 2; ++a) {
            if ((a * total) % n == 0 && dp[a].count((a * total) / n)) {
                return true;
            }
        }

        return false;
    }
};