Given an array A of integers, return the number of (contiguous, non-empty) subarrays that have a sum divisible by K.
Example 1:
Input: A = [4,5,0,-2,-3,1], K = 5
Output: 7
Explanation: There are 7 subarrays with a sum divisible by K = 5:
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]
Note:
1 <= A.length <= 30000
-10000 <= A[i] <= 10000
2 <= K <= 10000
定义一个sum[]数组,sum[i]=sum(A[0], A[1], ... A[i])。
对于每一个i,都从计算sum[i]~sum[A.length - 1]是否能够被K整除。
时间复杂度:O(n^2) 空间复杂度:O(n)
解法超时。本题规定了A[i]的大小,实际上sum还有溢出的风险。
public int subarraysDivByK(int[] A, int K) {
int[] sum = new int[A.length];
for (int i = 0; i < A.length; i++) {
sum[i] = (i == 0 ? 0 : sum[i - 1]) + A[i];
}
int ans = 0;
if (sum[0] % K == 0) ans++;
for (int i = 1; i < sum.length; i++) {
if (sum[i] % K == 0) ans++;
for (int j = i - 1; j >= 0; j--) {
int tmp = sum[i] - sum[j];
if (tmp % K == 0) {
ans++;
}
}
}
return ans;
}
其实上面的解法已经很接近了,但是思考一下为什么会超时?比如K=5,当sum[0]=4时,后面的sum[i]=5*n + 4都符合要求,即不需要判断每个和,只需要判断sum[i]%5的值就好了。
假设:
sum[i]=p*K+ri
sum[j]=q*K+rj
则:
sum[j]-sum[i]=(q-p)*K+rj-ri
若想sum[j]-sum[i]被K整除,只需要rj-ri能够被K整除即可,因为rj和ri都是对K的模,所以只要rj=ri即可。
所以这道题就变成了求对K的每个模的个数,然后对于每个0~i,分别两两组合,假设mod[i]=n,那么其i的个数为n*(n-1)/2.
时间复杂度:O(n) 空间复杂度:O(K)
public int subarraysDivByK(int[] A, int K) {
int mod[] = new int[K];
int cumSum = 0;
for (int i = 0; i < A.length; i++) {
cumSum += A[i];
// countSum % K 可能是负数,需要 +K
mod[((cumSum % K) + K) % K]++;
}
int ans = 0;
for (int i = 0; i < K; i++) {
if (mod[i] > 1) {
ans += (mod[i] * (mod[i] - 1)) / 2;
}
}
ans += mod[0];
return ans;
}