Split Array With Same Average
HardPrompt
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.
Example 1:
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:
Input: nums = [3,1]
Output: false
Constraints:
1 <= nums.length <= 300 <= nums[i] <= 104
Approaches
2 approaches with complexity analysis and trade-offs.
This is a straightforward approach that checks every possible partition of the array. We can use recursion to generate all subsets of nums to form the first array A. For each generated subset A, we check if it's a valid candidate (non-empty and not the full array). If it is, we calculate its average and the average of the remaining elements (forming array B). If the averages are equal, we've found a solution.
Algorithm
- The core idea is to check for each possible subset
Aifavg(A) == avg(nums). The conditionavg(A) == avg(B)simplifies toavg(A) == avg(nums). - This can be written as
sum(A) / len(A) == totalSum / N, orsum(A) * N == totalSum * len(A). - We can iterate through all possible lengths
kfor subsetAfrom1toN-1. - For each
k, we need to find if there's a subset of sizekwith a sum equal to(totalSum * k) / N. - A recursive helper function
find(index, k, targetSum)can be used to solve this subset sum problem. - The function explores two paths at each element
nums[index]: either include it in the subset or not. - The recursion stops when a subset of size
kis formed or when all elements are processed.
Walkthrough
The problem asks if we can partition nums into two non-empty arrays A and B with equal averages. This condition sum(A)/len(A) == sum(B)/len(B) is equivalent to sum(A)/len(A) == totalSum/N, where totalSum and N are the sum and length of the original array nums. This simplifies to sum(A) * N == totalSum * len(A).
This transforms the problem into a search for a non-empty proper subset A that satisfies this equation. The brute-force approach is to generate all possible subsets of nums, and for each subset, check if it satisfies the condition. We can implement this using a backtracking algorithm. We iterate through all possible lengths k for the subset A (from 1 to N/2) and for each k, we check if a subset of this size exists with the required sum (totalSum * k) / N. The check is done via a recursive function that explores including or excluding each element from the subset.
class Solution { public boolean splitArraySameAverage(int[] nums) { int n = nums.length; if (n < 2) { return false; } int totalSum = 0; for (int num : nums) { totalSum += num; } for (int k = 1; k <= n / 2; k++) { if ((totalSum * k) % n == 0) { int targetSum = (totalSum * k) / n; if (canFindSum(nums, 0, k, targetSum)) { return true; } } } return false; } private boolean canFindSum(int[] nums, int index, int k, int targetSum) { if (k == 0) { return targetSum == 0; } if (index >= nums.length || k < 0 || targetSum < 0) { return false; } // Option 1: Include nums[index] in the subset if (canFindSum(nums, index + 1, k - 1, targetSum - nums[index])) { return true; } // Option 2: Exclude nums[index] from the subset if (canFindSum(nums, index + 1, k, targetSum)) { return true; } return false; }}Complexity
Time
O(N * 2^N). The recursive function `canFindSum` explores `2^N` paths in the worst case. This function is called in a loop that runs up to `N/2` times.
Space
O(N), where N is the number of elements in `nums`. This is for the recursion stack depth.
Trade-offs
Pros
Simple to understand and implement.
It's a good starting point for thinking about the problem.
Cons
Highly inefficient due to its exponential time complexity.
Will result in a 'Time Limit Exceeded' (TLE) error on platforms like LeetCode for the given constraints.
This approach significantly improves upon the brute-force method by using a divide-and-conquer strategy known as "meet-in-the-middle". The key idea is to first transform the problem to a simpler one: finding a non-empty proper subset that sums to zero. Then, we split the array into two halves, generate all possible subset sums for each half, and then efficiently check if a sum from the left half and a sum from the right half can combine to zero.
Algorithm
- Problem Transformation: The condition
avg(A) == avg(nums)is equivalent tosum(A) * N == totalSum * len(A). Let's define a new arraybwhereb[i] = nums[i] * N - totalSum. The problem is now to find a non-empty proper subset ofbthat sums to 0. - Meet-in-the-Middle: Split the transformed array
binto two halves,leftandright. - Generate Sums: For each half, generate all possible subset sums for all possible subset sizes. Store this in
Set<Integer>[], wheresums[k]is a set of all sums achievable withkelements. - Combine Results: After generating sums for both halves (
left_sumsandright_sums), search for a pair of subsets (one from each half) that combine to a total sum of 0. Lets_lbe a sum fromleft_sumswith sizek_l, ands_rbe a sum fromright_sumswith sizek_r. We need to find a pair such thats_l + s_r = 0and1 <= k_l + k_r < N. - Iterate through all sums in the left half and check for the existence of its negation (
-s_l) in the right half's sums, ensuring the combined subset is a proper non-empty subset.
Walkthrough
The brute-force approach is too slow because N can be up to 30, making 2^30 operations infeasible. However, 2^15 is manageable. This suggests a meet-in-the-middle approach.
First, we transform the problem. The condition sum(A) / len(A) == totalSum / N is equivalent to N * sum(A) - len(A) * totalSum = 0. If we create a new array b where b[i] = nums[i] * N - totalSum, the problem becomes finding a non-empty proper subset of b that sums to 0.
Now, we apply meet-in-the-middle:
- Split
binto two halves,left(sizen1 = N/2) andright(sizen2 = N - n1). - For each half, we generate all possible subset sums for every possible subset size (
k). We store these inleft_sumsandright_sums, which are arrays of HashSets.left_sums[k]will contain all possible sums of subsets of sizekfrom theleftarray. - Finally, we try to combine sums from both halves. We iterate through each sum
s1of sizek1fromleft_sums. For eachs1, we search for a sums2of sizek2inright_sumssuch thats1 + s2 = 0. If we find such a pair, we must ensure the combined subset is a valid split, i.e., its sizek1 + k2is greater than 0 and less thanN. Since we iteratek1from 1, the non-empty condition is met. The sum of the entire transformed arraybis 0, so a combination with sizeNis trivial and not a valid split. Thus, we only need to check ifk1 + k2 < N.
import java.util.*; class Solution { public boolean splitArraySameAverage(int[] nums) { int n = nums.length; if (n < 2) return false; int totalSum = 0; for (int num : nums) totalSum += num; int[] b = new int[n]; for (int i = 0; i < n; i++) { b[i] = nums[i] * n - totalSum; } int n1 = n / 2; Set<Integer>[] leftSums = generateSums(b, 0, n1); Set<Integer>[] rightSums = generateSums(b, n1, n - n1); for (int k1 = 1; k1 <= n1; k1++) { for (int s1 : leftSums[k1]) { for (int k2 = 0; k2 <= n - n1; k2++) { if (rightSums[k2].contains(-s1)) { if (k1 + k2 > 0 && k1 + k2 < n) { return true; } } } } } return false; } private Set<Integer>[] generateSums(int[] arr, int start, int len) { Set<Integer>[] sums = new Set[len + 1]; for (int i = 0; i <= len; i++) { sums[i] = new HashSet<>(); } sums[0].add(0); for (int i = 0; i < len; i++) { int num = arr[start + i]; for (int k = i; k >= 0; k--) { for (int s : sums[k]) { sums[k + 1].add(s + num); } } } return sums; }}Complexity
Time
O(N^2 * 2^(N/2)). Generating the sums for both halves takes `O(N * 2^(N/2))`. The nested loops for checking the combination of sums take `O( (N/2) * 2^(N/2) * (N/2) )`, which simplifies to `O(N^2 * 2^(N/2))`. This is the dominant factor.
Space
O(N * 2^(N/2)). Each of the two `sums` arrays can store up to `2^(N/2)` sums in total across all subset sizes. The total number of `(k, s)` pairs is bounded by `N * 2^(N/2)`.
Trade-offs
Pros
Much more efficient than backtracking and is feasible for the given constraints.
Demonstrates an effective problem-solving technique for subset-related problems with moderately sized inputs.
Cons
More complex to understand and implement correctly compared to backtracking.
Requires careful handling of the problem transformation and the logic for combining results from the two halves.
Solutions
Solution
class Solution {public boolean splitArraySameAverage(int[] nums) { int n = nums.length; if (n == 1) { return false; } int s = Arrays.stream(nums).sum(); for (int i = 0; i < n; ++i) { nums[i] = nums[i] * n - s; } int m = n >> 1; Set<Integer> vis = new HashSet<>(); for (int i = 1; i < 1 << m; ++i) { int t = 0; for (int j = 0; j < m; ++j) { if (((i >> j) & 1) == 1) { t += nums[j]; } } if (t == 0) { return true; } vis.add(t); } for (int i = 1; i < 1 << (n - m); ++i) { int t = 0; for (int j = 0; j < (n - m); ++j) { if (((i >> j) & 1) == 1) { t += nums[m + j]; } } if (t == 0 || (i != (1 << (n - m)) - 1) && vis.contains(-t)) { return true; } } return false; }}class Solution {public: bool splitArraySameAverage(vector<int> &nums) { int n = nums.size(); if (n == 1) return false; int s = accumulate(nums.begin(), nums.end(), 0); for (int &v : nums) v = v * n - s; int m = n >> 1; unordered_set<int> vis; for (int i = 1; i < 1 << m; ++i) { int t = 0; for (int j = 0; j < m; ++j) if (i >> j & 1) t += nums[j]; if (t == 0) return true; vis.insert(t); } for (int i = 1; i < 1 << (n - m); ++i) { int t = 0; for (int j = 0; j < (n - m); ++j) if (i >> j & 1) t += nums[m + j]; if (t == 0 || (i != (1 << (n - m)) - 1 && vis.count(-t))) return true; } return false; }};class Solution: def splitArraySameAverage(self, nums: List[int]) -> bool: n = len(nums) if n == 1: return False s = sum(nums) for i, v in enumerate(nums): nums[i] = v * n - s m = n >> 1 vis = set() for i in range(1, 1 << m): t = sum(v for j, v in enumerate(nums[: m]) if i >> j & 1) if t == 0: return True vis . add(t) for i in range(1, 1 << (n - m)): t = sum(v for j, v in enumerate(nums[m:]) if i >> j & 1) if t == 0 or (i != (1 << (n - m)) - 1 and - t in vis): return True return FalseVideo walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.