Description
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
Solutions
Solution 1: Binary Search + Binary Enumeration
According to the problem requirements, we need to determine if the array nums can be divided into two subarrays A and B such that the average values of the two subarrays are equal.
Let the sum of the array nums be s, and the number of elements be n. The sum and number of elements of subarray A are s_1 and k, respectively. Then the sum of subarray B is s_2 = s - s_1, and the number of elements is n - k. Thus:
$$ \frac{s_1}{k} = \frac{s_2}{n - k} = \frac{s-s_1}{n-k} $$
Rearranging, we get:
$$ s_1 \times (n-k) = (s-s_1) \times k $$
Simplifying, we get:
$$ \frac{s_1}{k} = \frac{s}{n} $$
This means we need to find a subarray A such that its average value equals the average value of the array nums. We consider subtracting the average value of the array nums from each element of the array nums, transforming the problem into finding a subarray in the array nums whose sum is 0.
However, the average value of the array nums may not be an integer, and floating-point calculations may have precision issues. We can multiply each element of the array nums by n, i.e., nums[i] \leftarrow nums[i] × n. The above equation becomes:
$$ \frac{s_1\times n}{k} = s $$
Now, we subtract the integer s from each element of the array nums, transforming the problem into finding a subarray A in the array nums whose sum is 0.
The length of the array nums ranges from [1, 30]. If we use brute force to enumerate subarrays, the time complexity is O(2^n), which will time out. We can use binary search to reduce the time complexity to O(2^{n/2}).
We divide the array nums into left and right parts. The subarray A can exist in three cases:
- Subarray A is entirely in the left part of the array nums;
- Subarray A is entirely in the right part of the array nums;
- Subarray A is partially in the left part and partially in the right part of the array nums.
We can use binary enumeration to first enumerate the sums of all subarrays in the left part. If there is a subarray with a sum of 0, we return true immediately. Otherwise, we store the sums in a hash table vis. Then we enumerate the sums of all subarrays in the right part. If there is a subarray with a sum of 0, we return true immediately. Otherwise, we check if the hash table vis contains the opposite of the current sum. If it does, we return true.
Note that we cannot select all elements from both the left and right parts simultaneously, as this would leave subarray B empty, which does not meet the problem requirements. In implementation, we only need to consider n-1 elements of the array.
The time complexity is O(n × 2^{n⁄2}), and the space complexity is O(2^{n⁄2}).
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 False(code-box)
