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
Idea:
- See: https://leetcode.com/articles/subarray-sums-divisible-by-k/
- Calculate prefix sum of A, say P[i+1] = A[0] + ... + A[i], so subarray sum can be represented as P[j] - P[i] (j > i)
- Say C[i] = Count P[i] % K, i = 0..K-1
- Calculate combination sum for C[i] with i = 0..K-1
class Solution {
public:
int subarraysDivByK(vector<int>& A, int K) {
int n = A.size();
vector<int> prefix_sum(n + 1, 0);
for (int i = 0; i < n; i++)
prefix_sum[i + 1] = ((prefix_sum[i] + A[i]) % K + K) % K;
Note! Because prefix_sum can be negative, so here need ((..)% K + K) % K;- If prefix sum >= 0, the result is unchanged
- If prefix sum < 0, the result is complement value, e.g. P = -2, K = 3, ((P % K) + K) % K = 1
vector<int> freq(K, 0);
long long total = 0;
for (int i = 0; i <= n; i++) {
total += freq[prefix_sum[i]];
freq[prefix_sum[i]]++;
}
return total;
}
};
int subarraysDivByK(vector<int>& A, int K) {
const int n = A.size();
vector<int> prefsum(n + 1, 0);
for (int i = 0; i < n; ++i)
prefsum[i + 1] = prefsum[i] + A[i];
vector<int> count(K, 0); // 0 ~ K-1
vector<int> negCount(K+1, 0); // -K+1 ~ 0
for (int p : prefsum) {
int r = p % K;
if (r < 0)
++negCount[abs(r)];
else
++count[r];
}
int ans = 0;
for (int i = 0; i < K; ++i) {
int combine = count[i] + negCount[K - i];
ans += combine * (combine - 1) / 2;
}
return ans;
}