Problem

You are given an integer array nums.

You should move each element of nums into one of the two arrays A and B such that A and B are non-empty, and average(A) == average(B).

Return true if it is possible to achieve that and false otherwise.

Note that for an array arr, average(arr) is the sum of all the elements of arr over the length of arr.

Examples

Example 1

1
2
3

Input: nums = [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 an average of 4.5.

Example 2

1
2

Input: nums = [3,1]
Output: false

Constraints

  • 1 <= nums.length <= 30
  • 0 <= nums[i] <= 10^4

Solution

Method 1 – Subset Sum with Scaling and Bitmask DP

Intuition

To split the array into two non-empty groups with the same average, we need to find a non-empty subset whose sum is proportional to its size, matching the overall average. By scaling the problem to subset sums, we can use dynamic programming and bitmasking to efficiently check all possible subset sizes and sums.

Approach

  1. Let n be the length of nums and s be the total sum.
  2. For each possible subset size k from 1 to n//2:
    • The subset sum must be s * k / n and must be integer.
    • Use DP to check if any subset of size k has sum s * k / n.
  3. Use a set to store all possible sums for each subset size using bitmask DP.
  4. If any valid subset is found, return True; otherwise, return False.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

#include <vector>
#include <unordered_set>
using namespace std;
class Solution {
public:
    bool splitArraySameAverage(vector<int>& nums) {
        int n = nums.size(), s = 0;
        for (int x : nums) s += x;
        vector<unordered_set<int>> dp(n+1);
        dp[0].insert(0);
        for (int x : nums) {
            for (int k = n-1; k >= 0; --k) {
                for (int sum : dp[k]) {
                    dp[k+1].insert(sum + x);
                }
            }
        }
        for (int k = 1; k <= n/2; ++k) {
            if ((s * k) % n == 0 && dp[k].count(s * k / n)) return true;
        }
        return false;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

func splitArraySameAverage(nums []int) bool {
    n, s := len(nums), 0
    for _, x := range nums {
        s += x
    }
    dp := make([]map[int]struct{}, n+1)
    for i := range dp {
        dp[i] = make(map[int]struct{})
    }
    dp[0][0] = struct{}{}
    for _, x := range nums {
        for k := n - 1; k >= 0; k-- {
            for sum := range dp[k] {
                dp[k+1][sum+x] = struct{}{}
            }
        }
    }
    for k := 1; k <= n/2; k++ {
        if (s*k)%n == 0 {
            if _, ok := dp[k][s*k/n]; ok {
                return true
            }
        }
    }
    return false
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

import java.util.*;
class Solution {
    public boolean splitArraySameAverage(int[] nums) {
        int n = nums.length, s = 0;
        for (int x : nums) s += x;
        Set<Integer>[] dp = new Set[n+1];
        for (int i = 0; i <= n; ++i) dp[i] = new HashSet<>();
        dp[0].add(0);
        for (int x : nums) {
            for (int k = n-1; k >= 0; --k) {
                for (int sum : dp[k]) {
                    dp[k+1].add(sum + x);
                }
            }
        }
        for (int k = 1; k <= n/2; ++k) {
            if ((s * k) % n == 0 && dp[k].contains(s * k / n)) return true;
        }
        return false;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

class Solution {
    fun splitArraySameAverage(nums: IntArray): Boolean {
        val n = nums.size
        val s = nums.sum()
        val dp = Array(n+1) { mutableSetOf<Int>() }
        dp[0].add(0)
        for (x in nums) {
            for (k in n-1 downTo 0) {
                for (sum in dp[k]) {
                    dp[k+1].add(sum + x)
                }
            }
        }
        for (k in 1..n/2) {
            if ((s * k) % n == 0 && dp[k].contains(s * k / n)) return true
        }
        return false
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

class Solution:
    def splitArraySameAverage(self, nums: list[int]) -> bool:
        n, s = len(nums), sum(nums)
        dp = [set() for _ in range(n+1)]
        dp[0].add(0)
        for x in nums:
            for k in range(n-1, -1, -1):
                for total in dp[k]:
                    dp[k+1].add(total + x)
        for k in range(1, n//2+1):
            if (s * k) % n == 0 and (s * k // n) in dp[k]:
                return True
        return False
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

impl Solution {
    pub fn split_array_same_average(nums: Vec<i32>) -> bool {
        let n = nums.len();
        let s: i32 = nums.iter().sum();
        let mut dp = vec![std::collections::HashSet::new(); n+1];
        dp[0].insert(0);
        for &x in &nums {
            for k in (0..n).rev() {
                for &sum in dp[k].iter() {
                    dp[k+1].insert(sum + x);
                }
            }
        }
        for k in 1..=n/2 {
            if (s * k as i32) % n as i32 == 0 && dp[k].contains(&(s * k as i32 / n as i32)) {
                return true;
            }
        }
        false
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

class Solution {
    splitArraySameAverage(nums: number[]): boolean {
        const n = nums.length, s = nums.reduce((a, b) => a + b, 0);
        const dp: Set<number>[] = Array.from({length: n+1}, () => new Set());
        dp[0].add(0);
        for (const x of nums) {
            for (let k = n-1; k >= 0; --k) {
                for (const sum of dp[k]) {
                    dp[k+1].add(sum + x);
                }
            }
        }
        for (let k = 1; k <= n/2; ++k) {
            if ((s * k) % n === 0 && dp[k].has(s * k / n)) return true;
        }
        return false;
    }
}

Complexity

  • ⏰ Time complexity: O(n^2 * S) – n is the length, S is the sum of nums. For each subset size and sum.
  • 🧺 Space complexity: O(n * S) – For DP table storing possible sums.