805. Split Array With Same Average - Explanation

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 <= 30
0 <= nums[i] <= 10,000

Company Tags

Please upgrade to NeetCode Pro to view company tags.

1. Backtracking

class Solution:
    def splitArraySameAverage(self, nums: List[int]) -> bool:
        def backtrack(i, A, B):
            if i == len(nums):
                if not A or not B:
                    return False
                return sum(A) * len(B) == sum(B) * len(A)

            A.append(nums[i])
            if backtrack(i + 1, A, B):
                return True
            B.append(nums[i])
            A.pop()
            res = backtrack(i + 1, A, B)
            B.pop()
            return res

        return backtrack(0, [], [])

Time & Space Complexity

Time complexity: $O(n * 2 ^ n)$
Space complexity: $O(n)$

2. Memoization (Brute Force)

class Solution:
    def splitArraySameAverage(self, nums: List[int]) -> bool:
        total = sum(nums)
        n = len(nums)
        memo = {}

        def dfs(i, size, curr_sum):
            if (i, size, curr_sum) in memo:
                return memo[(i, size, curr_sum)]

            if size > 0 and size < n and curr_sum * (n - size) == (total - curr_sum) * size:
                return True

            if i == n:
                return False

            # include nums[i] in A
            if dfs(i + 1, size + 1, curr_sum + nums[i]):
                memo[(i, size, curr_sum)] = True
                return True

            # include nums[i] in B
            if dfs(i + 1, size, curr_sum):
                memo[(i, size, curr_sum)] = True
                return True

            memo[(i, size, curr_sum)] = False
            return False

        return dfs(0, 0, 0)

Time & Space Complexity

Time complexity: $O(n ^ 2 * s)$
Space complexity: $O(n ^ 2 * s)$

Where $n$ is the size of the input array $nums$ , and $s$ is the sum of the elements of the array.

3. Memoization (Optimal)

class Solution:
    def splitArraySameAverage(self, nums: List[int]) -> bool:
        n, total = len(nums), sum(nums)
        # len(A) = a, len(B) = b, let a <= b
        # avg(A) = avg(B)
        # sum(A) / a = sum(B) / b = sum(nums) / n
        # sum(A) / a = avg => sum(A) = a * avg
        # sum(A) = a * sum(nums) / n
        # Find if any subset exists with a * sum(nums) / n
        # a is in the range [1, (n//2)]

        memo = {}
        def dfs(i, a, s):
            if (i, a, s) in memo:
                return memo[(i, a, s)]
            if a == 0:
                return s == 0
            if i == n or a < 0:
                return False
            memo[(i, a, s)] = dfs(i + 1, a, s) or dfs(i + 1, a - 1, s - nums[i])
            return memo[(i, a, s)]

        for a in range(1, n // 2 + 1):
            if total * a % n == 0:
                if dfs(0, a, total * a // n):
                    return True

        return False

Time & Space Complexity

Time complexity: $O(n ^ 2 * s)$
Space complexity: $O(n ^ 2 * s)$

Where $n$ is the size of the input array $nums$ , and $s$ is the sum of the elements of the array.

4. Dynamic Programming (Bottom-Up)

class Solution:
    def splitArraySameAverage(self, nums: List[int]) -> bool:
        n = len(nums)
        # len(A) = a, len(B) = b, let a <= b
        # avg(A) = avg(B)
        # sum(A) / a = sum(B) / b = sum(nums) / n
        # sum(A) / a = avg => sum(A) = a * avg
        # sum(A) = a * sum(nums) / n
        # Find if any subset exists with a * sum(nums) / n
        # a is in the range [1, (n//2)]

        total = sum(nums)
        dp = [set() for _ in range(n // 2 + 1)]

        dp[0].add(0)
        for num in nums:
            for a in range(n // 2, 0, -1):
                for prev in dp[a - 1]:
                    dp[a].add(prev + num)

        for a in range(1, n // 2 + 1):
            if (a * total % n == 0) and (a * total // n) in dp[a]:
                return True

        return False

Time & Space Complexity

Time complexity: $O(n ^ 2 * s)$
Space complexity: $O(n ^ 2 * s)$

Where $n$ is the size of the input array $nums$ , and $s$ is the sum of the elements of the array.