805. Split Array With Same Average

https://leetcode.com/problems/split-array-with-same-average/

这题很有意思。和经典的动态规划将一个数组分隔成为两个总和相等的部分不同,它要求切成成为平均值相等的部分。

细想一个其实可以在原来问题上做改进,就是 `dp[x]` 存储的不仅仅是 “是否可以达到”,还要存储 “有几种不同个数” 可以达到。 比如 `dp[10] = [1,2,3]` 那么说明有 "有1个数字是10,有两个数字和也可以构成10,有三个数字也可以构成10". 当然这种值的方式可以改进,不使用列表,而是使用bits. 这也就是为什么题目里面限制 `A.length <= 30`.

一开始我还是按照传统思维去检查每个bits, 这样虽说时间复杂度上没有本质区别,但是还是出现超时。不过最后发现其实我们只需要将整个值向左移动一位就好了,没有循环就是高效。

def splitArraySameAverage(self, A: List[int]) -> bool:

    s = sum(A)
    n = len(A)
    from collections import defaultdict
    dp = defaultdict(int)
    dp[0] = 1

    for k, x in enumerate(A):
        for y, c in list(dp.items()):
            # 这里直接做位移就好了
            # c2 = 0
            # for i in range(n - 1):
            #     if (c & (1 << i)) == 0:
            #         continue
            #
            #     c2 |= (1 << (i + 1))
            #     if (x + y) * (n - i - 1) == (s - x - y) * (i + 1):
            #         return True
            #
            # if c2 != 0:
            #     dp[x + y] |= c2
            dp[x + y] |= (c << 1)

    # 最后我们枚举切分的长度就好
    for k in range(1, n - 1):
        # x / k == (s - x) / (n - k)
        # x = sk / n
        if (s * k) % n == 0:
            x = s * k // n
            if dp[x] & (1 << k):
                return True

    return False

使用我最开始的方法在Python下面是会超时的,但是Java就没有问题(276ms). 但是如果直接使用移位的话,那么Python可以通过(640ms), Java则可以到(15ms).

class Solution {
    public boolean splitArraySameAverage(int[] A) {
        int n = A.length;
        int sum = 0;
        for (int i = 0; i < n; i++) {
            sum += A[i];
        }
        int[] dp = new int[sum + 1];
        dp[0] = 1;

        for (int k = 0; k < n; k++) {
            int x = A[k];
            for (int y = sum - x; y >= 0; y--) {
                // 这里直接移位就好
                // int c = dp[y];
                // int c2 = 0;
                // for (int i = 0; i < (n - 1); i++) {
                //     if ((c & (1 << i)) == 0)
                //         continue;

                //     c2 |= (1 << (i + 1));
                //     if (((x + y) * (n - i - 1)) == ((sum - x - y) * (i + 1))) {
                //         return true;
                //     }
                // }
                // dp[x + y] |= c2;
                dp[x + y] |= (dp[y] << 1);
            }
        }

        for (int k = 1; k < n; k++) {
            if ((sum * k) % n == 0) {
                int x = sum * k / n;
                if ((dp[x] & (1 << k)) != 0) {
                    return true;
                }
            }
        }
        return false;
    }
}