题目描述:
给定一个整数数组 A
,返回其中元素之和可被 K
整除的(连续、非空)子数组的数目。
题解:
解题的大致思路和leetcode 570类似。子数组$(i,j)$和可以用$j$的前缀和减去$i$的前缀和表示,我们考虑所有以$end$结尾的子数组中,有多少个是符合题意的。在以$end$结尾的子数组中,假设v1为$end$对应的前缀和,v2为子数组起始位置$start$对应的前缀和。如果$start$为满足题意的位置,那么就有$(v2 - v1) mod k = 0$,也就是说对于每个$end$,我们想知道有多少个$start$使得$(v2 - v1) mod k = 0$。既然要计算的是次数,自然想到用map做一个映射,这样就不用去遍历所有的$start$。但是$(v2 - v1) mod k = 0$这个条件并不好用map表示,我们换一种表现形式:$ v2%K - v1%K == 0 or v2%K - v1%K == K$,这个是根据同余定理转换的。有了这么一个转换之后,我们只需要用map记录每个前缀和对K取余之后的值出现的次数,就可得出以$end$结尾的子数组中有多少个符合题意的结果,然后遍历所有可能的$end$就可以得到最终答案了。代码如下所示:
class Solution { public: int subarraysDivByK(vector<int>& A, int K) { unordered_map<int,int> mp; int len = A.size(); int ans = 0; int sum = 0; mp[0] = 1; for(int i=0;i<len;i++) { sum += A[i]; // cout << -9%9 <<endl; if(mp.find(sum%K) != mp.end()) ans += mp[sum%K]; int tmp = (sum%K > 0)?-K:K; if(mp.find(sum%K + tmp) != mp.end()) ans += mp[sum%K + tmp]; mp[sum%K]++; } // if(abs(sum)%K == 0) ans++; return ans; } };