805. Split Array With Same Average


Problem Description

The problem presents us with an integer array called nums. Our goal is to determine whether we can divide this array into two non-empty arrays, A and B, such that the average (or mean) of the elements in A is equal to the average of elements in B. This means the sum of all elements in A divided by the number of elements in A should be equal to the sum of all elements in B divided by the number of elements in B. We need to return true if such a split is possible, and false otherwise.

Intuition

Facing this problem, the immediate intuition might be to try all possible splits of the array nums into two arrays and check their averages. However, this approach would be inefficient due to the high number of possible splits.

The insight required to arrive at a more efficient solution involves understanding that if two arrays have the same average, then the sum of all elements in both arrays combined, multiplied by the length of one of the arrays must equal the sum of the elements in that array multiplied by the length of both arrays combined. We can use this property to transform the elements of the original array.

Thus, the solution approach involves pre-processing of the input array. All elements v are transformed into v*n - s where n is the length of the input array and s is the sum of all elements of the array. This transformation simplifies the comparison of sums and counters the different lengths of arrays A and B.

After this transformation, a bit manipulation technique is used to generate all possible subset sums from both halves of the transformed array using binary representations of numbers from 1 up to 2^m - 1. These generated sums are stored or compared to zero or to the negative of sums from the other half, since the aim is to find a zero-sum subset, which would indicate equal averages in accordance with the transformation applied.

The core of the algorithm checks for a zero-sum subset among all possible subsets of each half. If such subset is found in either half or the negation of a sum from half is found in the other, we can guarantee that there exists a configuration of A and B with equal averages.

Learn more about Math, Dynamic Programming and Bitmask patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece:

What is the best way of checking if an element exists in a sorted array once in terms of time complexity? Select the best that applies.

Solution Approach

The solution approach makes use of several important concepts and data structures:

1. Transformation of Elements:

  • As part of the preprocessing, every element v in the array is replaced with v*n - s, where n is the total number of elements in nums and s is their sum. This transformation balances the scale of the sums when considering different lengths for A and B.

2. Bit Manipulation for Subset Sums:

  • The idea is to use the bits of a number to represent which elements are included in a subset. If nums has n elements, for every integer i from 1 to 2^n - 1, the binary representation of i indicates which elements to include (bit 1) or exclude (bit 0) from a subset.
  • For example, if nums has n = 3 and the current i = 3 (binary 011), it means include nums[1] and nums[2] in the subset, but exclude nums[0].

3. Use of Set for Storing Intermediate Sums:

  • A set vis is used to store sums of subsets for the first half of the array.
  • The choice of a set structure allows for quick lookup to check if the inverse of a sum from the second half exists.

4. Dealing with Different Subset Sizes:

  • Since trying all possible subset sizes would be highly inefficient, the algorithm cleverly avoids full combinatorial explosion by dividing the array into two halves and generating subsets from them independently.
  • This division reduces the number of subsets that need to be considered. A full combination is never actually generated; rather, only potential matching combinations between halves are evaluated on-the-fly.

5. Checking for Zero Sums or Matching Negative Sums:

  • The algorithm uses sums of subsets and checks if any of them equal 0. A zero sum indicates that a valid split can be formed because the transformation ensures that a zero sum corresponds to equal averages.
  • Additionally, if the current sum t is not zero, the algorithm checks if -t (the negative value of t) is in the set of sums from the first half (for subsets of the second half). This indicates that the part of the first half adds up to A, and the part of the second half adds up to B, and both A and B give the same average by counterbalancing each other.

6. Edge Cases:

  • If n == 1, it is impossible to split nums into two non-empty arrays. The function immediately returns false.

To perform the computations efficiently, bitwise operations are employed. When the loop variable i is iterated from 1 to 1 << m, where m is the max index for the loop (half the length of the nums array), it generates all possible subsets for the corresponding half. The expression i >> j & 1 checks if the j-th element should be included in the subset represented by i (i.e., whether the j-th bit of i is a 1).

Combining all these concepts, the algorithm efficiently checks if there exists a valid split where both halves have the same average, which could otherwise be a computationally expensive problem if a brute force method were to be used.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

How does quick sort divide the problem into subproblems?

Example Walkthrough

Consider the array nums = [1, 2, 3, 4, 5, 6]. The sum of this array is s = 21 and it has n = 6 elements. The goal is to see if we can split nums into two non-empty arrays, A and B such that the average of A is equal to the average of B.

1. Transformation of Elements: We transform each element v into v*n - s which in this case equates to:

  • 1*6 - 21 = -15
  • 2*6 - 21 = -9
  • 3*6 - 21 = -3
  • 4*6 - 21 = 3
  • 5*6 - 21 = 9
  • 6*6 - 21 = 15

The transformed array now looks like this: [-15, -9, -3, 3, 9, 15].

2. Bit Manipulation for Subset Sums: Let's consider generating subsets for the first three elements (first half). We look at integers from 1 to 2^3 - 1 (which is 7):

  1. For i = 1 (binary 001), include nums[2]: subset is [-3].
  2. For i = 2 (binary 010), include nums[1]: subset is [-9].
  3. For i = 3 (binary 011), include nums[1] and nums[2]: subset is [-12]. ... and so forth until i = 7 (binary 111), which includes all first half elements: subset is [-27].

3. Use of Set for Storing Intermediate Sums: The set vis stores these sums: {-3, -9, -12, ..., -27}.

4. Dealing with Different Subset Sizes: Similar subsets are generated for the next three elements (second half) and stored in another set.

5. Checking for Zero Sums or Matching Negative Sums: As we look for sums of subsets in the second half, we check if the sum t equals 0 or if -t exists in vis. For example, let's say we find a subset in the second half that totals -9. Checking our vis set reveals a 9 which indicates that when these two subsets are combined, their total is 0. This means their non-transformed elements can be split into two arrays A and B with the same average.

6. Edge Cases: Since n is greater than 1, we proceed with the calculation.

Putting this into practice for our example, imagine we find a subset in the second half that has a sum of 9 (such as [5] when not transformed), and in vis we find -9 (from [-9] which is [2] when not transformed). The transformation assures us that the average of [5] matches the average of [2]. Hence, we can return true.

Even though the example only shows a single element from each half for brevity, in the actual approach, all subsets would be considered, and a comprehensive check would ensure the right pair of subsets is found.

Not Sure What to Study? Take the 2-min Quiz:

Which of the following shows the order of node visit in a Breadth-first Search?

Python Solution

1class Solution:
2    def splitArraySameAverage(self, nums: List[int]) -> bool:
3        # Get the number of elements in the array
4        num_elements = len(nums)
5      
6        # Check base condition. If only one element, can't split
7        if num_elements == 1:
8            return False
9      
10        # Calculate the sum of all the elements
11        total_sum = sum(nums)
12      
13        # Adjust each number in the array to be the difference between
14        # the number times the number of elements and the total sum
15        for i, value in enumerate(nums):
16            nums[i] = value * num_elements - total_sum
17      
18        # Let m be half of num_elements, rounded down
19        half_num_elements = num_elements >> 1
20      
21        # Initialize a set to keep track of seen sums in the first half
22        seen_sums_first_half = set()
23      
24        # Iterate over all possible subsets of the first half of the array
25        for bitmask in range(1, 1 << half_num_elements):
26            # Calculate the subset sum of current bitmask in the first half
27            subset_sum = sum(value for j, value in enumerate(nums[:half_num_elements]) if bitmask >> j & 1)
28            # If subset sum is 0, a split is found
29            if subset_sum == 0:
30                return True
31            # Store this subset sum as seen
32            seen_sums_first_half.add(subset_sum)
33      
34        # Iterate over all possible subsets of the second half of the array
35        for bitmask in range(1, 1 << (num_elements - half_num_elements)):
36            # Calculate the subset sum of current bitmask in the second half
37            subset_sum = sum(value for j, value in enumerate(nums[half_num_elements:]) if bitmask >> j & 1)
38          
39            # If subset sum is 0 or the negative of this subset sum was seen in the first half, a split is found
40            # In the first loop, this check prevents considering the entire second half as a subset.
41            if subset_sum == 0 or (bitmask != (1 << (num_elements - half_num_elements)) - 1 and -subset_sum in seen_sums_first_half):
42                return True
43      
44        # If no split is found where both parts have the same average, return False
45        return False
46

Java Solution

1class Solution {
2    public boolean splitArraySameAverage(int[] nums) {
3        int length = nums.length;
4        // An array with only one element can't be split to satisfy the condition
5        if (length == 1) {
6            return false;
7        }
8
9        // Calculate the sum of all array elements
10        int sum = Arrays.stream(nums).sum();
11
12        // Adjust each number in the array to be the difference between
13        // the number * length and the total sum. This transformation
14        // simplifies the average comparison to a sum comparison
15        for (int i = 0; i < length; ++i) {
16            nums[i] = nums[i] * length - sum;
17        }
18
19        // Calculate half the length of the array
20        int halfLength = length >> 1;
21
22        // Initialize a set to keep track of sums that we've visited
23        Set<Integer> seenSums = new HashSet<>();
24
25        // Evaluate all possible sums for the first half of the array
26        for (int i = 1; i < 1 << halfLength; ++i) {
27            int currentSum = 0;
28            for (int j = 0; j < halfLength; ++j) {
29                // Check if the j-th bit of i is set and, if so, add the number at index j to currentSum
30                if (((i >> j) & 1) == 1) {
31                    currentSum += nums[j];
32                }
33            }
34
35            // If a zero sum is found, then it's possible to split the array
36            if (currentSum == 0) {
37                return true;
38            }
39
40            // Add the calculated sum to the set
41            seenSums.add(currentSum);
42        }
43
44        // Now, evaluate all possible sums for the second half of the array
45        for (int i = 1; i < 1 << (length - halfLength); ++i) {
46            int currentSum = 0;
47            for (int j = 0; j < (length - halfLength); ++j) {
48                // Check if the j-th bit of i is set and, if so, add the number at index halfLength+j to currentSum
49                if (((i >> j) & 1) == 1) {
50                    currentSum += nums[halfLength + j];
51                }
52            }
53
54            // A sum of zero means we found a potential split, and if the negated sum was seen in the first half, a split is possible
55            if (currentSum == 0 || (i != (1 << (length - halfLength)) - 1) && seenSums.contains(-currentSum)) {
56                return true;
57            }
58        }
59
60        // No split found that satisfies the condition
61        return false;
62    }
63}
64

C++ Solution

1#include <vector>
2#include <numeric>
3#include <unordered_set>
4
5class Solution {
6public:
7    bool splitArraySameAverage(std::vector<int>& nums) {
8        int size = nums.size();
9
10        // If there's only one element, you can't split it into two non-empty parts.
11        if (size == 1) return false;
12
13        // Calculate the total sum of all elements in the array.
14        int totalSum = std::accumulate(nums.begin(), nums.end(), 0);
15
16        // Normalize each number by the formula: v * size - totalSum.
17        for (int& value : nums) value = value * size - totalSum;
18
19        // m is half the size of the array, for splitting the array into two parts.
20        int halfSize = size >> 1;
21
22        // Use a set to store sums of all possible combinations in the first half.
23        std::unordered_set<int> sumsFirstHalf;
24
25        // Check all possible subsets in the first half of the array.
26        for (int i = 1; i < 1 << halfSize; ++i) {
27            int currentSum = 0;
28
29            // Calculate the sum of the current subset.
30            for (int j = 0; j < halfSize; ++j)
31                if (i >> j & 1) currentSum += nums[j];
32
33            // If a zero-sum subset is found, the answer is true.
34            if (currentSum == 0) return true;
35
36            // Insert current sum into the set.
37            sumsFirstHalf.insert(currentSum);
38        }
39
40        // Check all possible subsets in the second half of the array.
41        for (int i = 1; i < 1 << (size - halfSize); ++i) {
42            int currentSum = 0;
43
44            // Calculate the sum of the current subset.
45            for (int j = 0; j < (size - halfSize); ++j)
46                if (i >> j & 1) currentSum += nums[halfSize + j];
47
48            // Check if either of the current subset sums to zero or they add up to a zero sum
49            // with a subset sum from the other half. If so, we can split the array satisfying the condition.
50            if (currentSum == 0 || (i != (1 << (size - halfSize)) - 1 && sumsFirstHalf.count(-currentSum))) return true;
51        }
52
53        // If none of the subsets satisfy the criteria, return false.
54        return false;
55    }
56};
57

Typescript Solution

1function splitArraySameAverage(nums: number[]): boolean {
2    let size = nums.length;
3
4    // If there's only one element, you can't split it into two non-empty parts
5    if (size === 1) {
6        return false;
7    }
8
9    // Calculate the total sum of all elements in the array
10    let totalSum = nums.reduce((accumulator, currentValue) => accumulator + currentValue);
11
12    // Normalize each number by the formula: v * size - totalSum
13    nums = nums.map(value => value * size - totalSum);
14
15    // 'halfSize' is half the size of the array, for splitting the array into two parts
16    let halfSize = Math.floor(size / 2);
17
18    // Use a set to store sums of all possible combinations in the first half
19    let sumsFirstHalf = new Set<number>();
20
21    // Function to compute the sums of all subsets of a given array
22    const computeSubsetSums = (arr: number[]): Set<number> => {
23        let subsetSums = new Set<number>();
24        let subsetsCount = 1 << arr.length; // Total subsets for the given array
25
26        for (let i = 1; i < subsetsCount; i++) {
27            let currentSum = 0;
28          
29            for (let j = 0; j < arr.length; j++) {
30                if ((i >> j) & 1) {
31                    currentSum += arr[j];
32                }
33            }
34
35            if (currentSum === 0) {
36                return new Set([0]);
37            }
38          
39            subsetSums.add(currentSum);
40        }
41      
42        return subsetSums;
43    };
44
45    // Compute the sums of all possible subsets for the first half of the array
46    sumsFirstHalf = computeSubsetSums(nums.slice(0, halfSize));
47
48    if (sumsFirstHalf.has(0)) {
49        return true;
50    }
51
52    // Compute the sums of all possible subsets for the second half of the array
53    let sumsSecondHalf = computeSubsetSums(nums.slice(halfSize));
54
55    // Check if either of the current subset sums to zero or they add up to a zero sum
56    // with a subset sum from the other half. If so, we can split the array satisfying the condition.
57    for (let sum of sumsSecondHalf) {
58        if (sum === 0 || sumsFirstHalf.has(-sum)) {
59            return true;
60        }
61    }
62
63    // If none of the subsets satisfy the criteria, return false
64    return false;
65}
66
Fast Track Your Learning with Our Quick Skills Quiz:

How many times is a tree node visited in a depth first search?

Time and Space Complexity

The given Python code aims to find out if an array can be split into two non-empty parts with the same average.

Time Complexity:

The time complexity of the algorithm is dominated by the subset sum computations done by the two for loops:

  1. The first for loop iterates over all subsets of the first half of the array, which leads to a complexity of O(2^(n/2)), as there are 2^(n/2) possible subsets.

  2. The second for loop iterates over all subsets of the second half of the array. The time complexity here is also O(2^(n/2)), for the same reason.

Since these two loops are sequential and not nested, we sum their complexities, giving us a total time complexity of O(2^(n/2) + 2^(n/2)), which simplifies to O(2^(n/2)).

Space Complexity:

The space complexity is the sum of the space required to store the modified nums array and the vis set that contains subset sums of the first half of the array:

  • The vis set can hold at most 2^(n/2) values in the worst case, which gives us a space complexity of O(2^(n/2)) for the set.

  • The nums array is modified in-place and does not require additional space proportional to the size of nums.

Therefore, the total space complexity of the algorithm is O(2^(n/2)).

Learn more about how to find time and space complexity quickly.


Recommended Readings


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫