Leetcode：805 Split Array With Same Average

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

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.

分析

这道题目有点难，因为长度为 N 的数组它的组合有 2**N 种。
题目给出数字范围 N < 30 且每个数均小于 104。
30 * 104 约等于 3200，只需要求 dp[使用的数字个数][数字和] 即可
此外，本地和在一个数组中找出和等于某个数的方法有点类似

实现

class Solution(object):
    def splitArraySameAverage(self, nums):        
        N, S = len(nums), sum(nums)
        H = N //2 + 1
        if 2 > N:
            return False
        if N == 2:
            return nums[0] == nums[1]
        
        divid = False

        for i in range(1, H):  ## 这里是优化速度的，速度速度从 1348 ms 提升到 80 ms
            if S * i % N == 0:
                divid = True
                break
            
        if not divid:
            return False
        
        dp =  [[{} for _ in range(H)] for _ in range(2)]
        dp[1][1][nums[0]] = 1
        dp[1][0][0] = 1
        
        for i in range(2, N+1):

            now_level, pre_level = i % 2, (i-1) % 2
            dp[now_level] = [{} for _ in range(H)]

            for used in range(H):
                dp[now_level][used].update(dp[pre_level][used])
                for j in dp[pre_level][used]:
                    n_j = j+nums[i-1]
                    if N * n_j == (used+1)*S :
                        return True
                    if H > used+1:
                        dp[now_level][used+1][n_j] = 1

        return False

总结

提交很多次，大部分是 index out 导致的 runtime error。部分是 TLE 
第一次通过的版本速度有点慢，只 beat 25% 的提交。优化后的速度 beat 75 % 的速度