LC 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; } }