Want Help Cracking FAANG?

(Then click this)

×
Back to Question Bank

805. Split Array With Same Average - Leetcode Solution

Problem Description

Given an array of integers A, your task is to determine if it is possible to split the array into two non-empty subsets B and C such that the average of both subsets is the same. In other words, you want to find at least one non-empty subset B (not equal to A itself) so that:

  • average(B) == average(C)
  • B and C are disjoint and together form A
  • Each subset must have at least one element, and you cannot reuse elements

For example, if A = [1,2,3,4,5,6,7,8], you can split into B = [1,4,5,8] and C = [2,3,6,7] since both have an average of 4.5.

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

Thought Process

At first glance, the problem seems to require checking all possible subsets of A to see if any can be split off such that both parts have the same average. However, with up to 30 elements, checking every possible subset directly is computationally infeasible (since there are 2n subsets).

The key insight is to realize that the average constraint translates into a sum and size constraint for possible subsets. If the whole array has sum S and length N, then for a subset of size k to have the same average as the whole array, it must have sum sumB = S * k / N. This means we only need to check if there exists a non-empty subset of size k (where 1 ≤ k ≤ N//2) whose sum is exactly sumB (and sumB is integer).

Brute force would still be too slow, so we need a way to efficiently check for such subsets.

Solution Approach

To efficiently solve the problem, we use the following strategy:

  1. Calculate total sum and length: Let S be the sum of A, and N its length.
  2. Iterate over possible subset sizes: For each k from 1 to N//2:
    • Check if S * k % N == 0. If not, skip, since sum must be integer.
    • Let target = S * k // N. We need to check if any subset of size k sums to target.
  3. Use dynamic programming (DP) to check subset sums:
    • We can use a DP array (or dictionary) where dp[size] is the set of all possible sums using size elements from A.
    • Initialize dp[0] = {0}.
    • For each number in A, update dp for increasing sizes.
  4. Check if target is in dp[k]: If so, return True. If no such subset is found for any k, return False.

This approach is efficient because:

  • We avoid redundant checks by using DP.
  • We only consider subset sizes that can possibly match the average.
  • We stop early if a valid subset is found.

Example Walkthrough

Let's walk through the example A = [1,2,3,4,5,6,7,8]:

  • Total sum: S = 36, Length: N = 8, Average: 4.5
  • For k = 1: 36 * 1 / 8 = 4.5 (not integer) → skip
  • For k = 2: 36 * 2 / 8 = 9 (integer)
  • Check if any subset of size 2 sums to 9: possible (e.g. [1,8], [2,7], [3,6], [4,5])
  • But these subsets leave 6 elements, whose sum is 27, average is 4.5 (same as original), so valid!
  • Therefore, the answer is True.

The DP builds up all possible sums for each subset size, and quickly finds that a subset of size 2 with sum 9 exists.

Time and Space Complexity

  • Brute-force approach: Would require checking all possible subsets, leading to O(2N) time and space, which is intractable for N up to 30.
  • Optimized DP approach:
    • There are up to N//2 subset sizes to check.
    • For each size, we maintain a set of possible sums. The number of possible sums is bounded by the sum of the array and the number of elements, so it's manageable for N ≤ 30.
    • Overall, the time complexity is O(N × S × N), where S is the average sum per subset size. This is feasible for the allowed constraints.
    • Space complexity is O(N × S) for the DP storage.

Summary

The key insight is that the "same average" constraint can be translated into a subset sum problem with a specific target, which allows us to use dynamic programming to efficiently check all possible subset sizes. By focusing only on subset sizes and sums that could possibly match the average, and using DP to avoid redundant calculations, we can solve the problem efficiently even for arrays with up to 30 elements.

Code Implementation

from collections import defaultdict

class Solution:
    def splitArraySameAverage(self, A):
        N = len(A)
        S = sum(A)
        # Early pruning: impossible if no k in 1..N//2 with integer sum
        possible = False
        for k in range(1, N//2 + 1):
            if S * k % N == 0:
                possible = True
                break
        if not possible:
            return False

        dp = [set() for _ in range(N+1)]
        dp[0].add(0)
        for num in A:
            # go backwards to avoid reusing num in the same dp update
            for size in range(N-1, 0, -1):
                for prev_sum in dp[size-1]:
                    dp[size].add(prev_sum + num)
        for k in range(1, N//2 + 1):
            if S * k % N != 0:
                continue
            target = S * k // N
            if target in dp[k]:
                return True
        return False
      
#include <vector>
#include <numeric>
#include <unordered_set>
using namespace std;

class Solution {
public:
    bool splitArraySameAverage(vector<int>& A) {
        int N = A.size(), S = accumulate(A.begin(), A.end(), 0);
        bool possible = false;
        for (int k = 1; k <= N/2; ++k) {
            if ((S * k) % N == 0) {
                possible = true;
                break;
            }
        }
        if (!possible) return false;

        vector<unordered_set<int>> dp(N+1);
        dp[0].insert(0);
        for (int num : A) {
            for (int size = N-1; size >= 1; --size) {
                for (int prev_sum : dp[size-1]) {
                    dp[size].insert(prev_sum + num);
                }
            }
        }
        for (int k = 1; k <= N/2; ++k) {
            if ((S * k) % N != 0) continue;
            int target = S * k / N;
            if (dp[k].count(target)) return true;
        }
        return false;
    }
};
      
import java.util.*;

class Solution {
    public boolean splitArraySameAverage(int[] A) {
        int N = A.length, S = 0;
        for (int num : A) S += num;
        boolean possible = false;
        for (int k = 1; k <= N / 2; ++k) {
            if ((S * k) % N == 0) {
                possible = true;
                break;
            }
        }
        if (!possible) return false;

        HashSet<Integer>[] dp = new HashSet[N + 1];
        for (int i = 0; i <= N; ++i) dp[i] = new HashSet<>();
        dp[0].add(0);

        for (int num : A) {
            for (int size = N - 1; size >= 1; --size) {
                for (int prevSum : dp[size - 1]) {
                    dp[size].add(prevSum + num);
                }
            }
        }

        for (int k = 1; k <= N / 2; ++k) {
            if ((S * k) % N != 0) continue;
            int target = S * k / N;
            if (dp[k].contains(target)) return true;
        }
        return false;
    }
}
      
var splitArraySameAverage = function(A) {
    const N = A.length;
    const S = A.reduce((a, b) => a + b, 0);
    let possible = false;
    for (let k = 1; k <= Math.floor(N / 2); ++k) {
        if ((S * k) % N === 0) {
            possible = true;
            break;
        }
    }
    if (!possible) return false;

    const dp = Array.from({length: N + 1}, () => new Set());
    dp[0].add(0);
    for (const num of A) {
        for (let size = N - 1; size >= 1; --size) {
            for (const prevSum of dp[size - 1]) {
                dp[size].add(prevSum + num);
            }
        }
    }
    for (let k = 1; k <= Math.floor(N / 2); ++k) {
        if ((S * k) % N !== 0) continue;
        const target = S * k / N;
        if (dp[k].has(target)) return true;
    }
    return false;
};