<p>You are given an integer array <code>nums</code>.</p>

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

<p>Return <code>true</code> if it is possible to achieve that and <code>false</code> otherwise.</p>

<p><strong>Note</strong> that for an array <code>arr</code>, <code>average(arr)</code> is the sum of all the elements of <code>arr</code> over the length of <code>arr</code>.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<pre>
<strong>Input:</strong> nums = [1,2,3,4,5,6,7,8]
<strong>Output:</strong> true
<strong>Explanation:</strong> 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.
</pre>

<p><strong class="example">Example 2:</strong></p>

<pre>
<strong>Input:</strong> nums = [3,1]
<strong>Output:</strong> false
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= nums.length &lt;= 30</code></li>
	<li><code>0 &lt;= nums[i] &lt;= 10<sup>4</sup></code></li>
</ul>

from functools import lru_cache

class Solution:
    def splitArraySameAverage(self, nums):
        """
        :type nums: List[int]
        :rtype: bool
        """
        
        # Thinking that this will be O(n^2)

        # Brute force: Split point based on if you include element i in A or if
        # you include element i in B
        """
        def dfs(i, A, B):
            if i >= len(nums):
                if not A or not B:
                    return False

                return (sum(A) / len(A)) == (sum(B) / len(B))
            
            # Let's say you add i into A
            A.append(nums[i])
            left = dfs(i + 1, A, B)
            A.remove(nums[i])

            B.append(nums[i])
            right = dfs(i + 1, A, B)
            B.remove(nums[i])

            return left or right

        return dfs(0, [], [])
        """

        # Do a DP approach based on this
        total_sum = sum(nums)
        n = len(nums)

        # Can you construct a subset of length curr_len with sum curr_sum using
        # elements from i onwards?
        @lru_cache(None)
        def dp(i, curr_len, curr_sum):
            if curr_len == 0:
                return curr_sum == 0
            if i >= n or curr_len < 0 or curr_sum < 0:
                return False

            # Include nums[i] in subset
            if dp(i + 1, curr_len - 1, curr_sum - nums[i]):
                return True
            
            # Exclude nums[i] in subset
            return dp(i + 1, curr_len, curr_sum)
        
        # Try to find any subset with the same average as the entire array
        for length in range(1, n):
            if total_sum * length % n == 0:
                target_sum = total_sum * length // n
                if dp(0, length, target_sum):
                    return True
        
        return False

        






        
90220