题目描述 和可被 K 整除的子数组
给定一个整数数组 A,返回其中元素之和可被 K 整除的(连续、非空)子数组的数目。
示例
输入:A = [4,5,0,-2,-3,1], K = 5
输出:7
解释:
有 7 个子数组满足其元素之和可被 K = 5 整除:
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]
解题思路
第一反应是双层循环,超时,gg。想了半天也不会,只好看看大佬们怎么写的。啊,原来是用前缀和啊,嗯,还是不懂。那好吧,只能再看看大佬们是怎么解释的。一菜毁所有!
参考
我们首先定义一个前缀和数组 prefix,prefix[0]=0,prefix[1]=A[0],prefix[2]=prefix[1]+A[1]……总之,prefix就是从第0个元素开始加和的结果。
以 A=[4,5,0,-2,-3,1],K=5 为例,然后我们得到下表:
其中 prefix%K 就是 prefix 对 5 求余的结果。
诀窍就在这里,我们观察 index 1-3 中 prefix%5 出现的3个4:
prefix[1] = ΣA[0], prefix[2] = ΣA[0-1]
而ΣA[0]%5 = ΣA[0-1]%5,这就表明Σ0和Σ0-1的差值必然是5的倍数,得到 ΣA[0-1] - ΣA[0] → A[1]%5==0 → {5}是一个符合集合。
同理,在第1个4和第3个4之间:
prefix[1] = ΣA[0], prefix[3] = ΣA[0-2]
→ ΣA[0-2] - ΣA[0] → ΣA[1-2]%5==0 → {5,0}是一个符合集合。
同理,在第2个4和第3个4之间,夹了一个A[2]=0:
→ {0}是一个符合集合。
这样,我们可以知道,当同样的数字出现2次时,表明1个符合集合出现了;当同样的数字出现3次时,表明3个符合集合出现了。这说明在4出现的3次中,两两一对,共有3对代表了3个符合的集合。
设出现4的次数为n,则符合的集合数量为C(n,2),
C(n,2)=n!/((n-2)!2)=n(n-1)/2。
在上面的例子中:
0出现了2次,result加上1;
3出现了1次,不能配对,所以result不变;
4出现了4次,result加上C(4,2)=6;
最终符合的集合数量为7。
这里我们用数组 cnt 来记录余数为 0-4 的次数。
至于 (n%K+K)%K 的目的在于将余数为负数的转化为正数,例如 n%5 in {-4,-3,-2,-1,0,1,2,3,4},则 (n%5+5)%5 in {0,1,2,3,4}。
代码
class Solution {
public:
int subarraysDivByK(vector<int>& A, int K) {
vector<int> remain(K,0);
int sum=0;
remain[0]=1; //因为一开始前0个数和为0
for(auto x: A){
sum+=x;
++remain[(sum%K+K)%K];
}
int ans=0;
for(auto x: remain){
if(x!=0&&x!=1) ans+=(x)*(x-1)/2;
}
return ans;
}
};