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