leetcode 974. 和可被 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]

提示：
1 <= A.length <= 30000
-10000 <= A[i] <= 10000
2 <= K <= 10000

1 暴力解法

开始拿到题，第一反应是两层循环，外层 i 从 0-N，内层 j 从 i 到 N，外层里用一个sum来求和，用sum%K==0来判断能否整除。这样时间复杂度是O(n^2)，果然超时了。

看了别人的答案，依然理解不了，之后才知道要用前缀和的方法来达到O(n)的复杂度。

2 前缀和解法

我们首先定义一个前缀和数组 prefix，prefix[0]=0，prefix[1]=A[0]，prefix[2]=prefix[1]+A[1]……总之，prefix就是从第0个元素开始加和的结果。

我们这样得到 prefix ：

int subarraysDivByK(vector<int>& A, int K) {
	int N = A.size();
    vector<int> prefix(N+1);
    for (int i=0;i<N;++i)
    {
        prefix[i+1]=prefix[i]+A[i];
    }

以 A=[4,5,0,-2,-3,1]，K=5 为例，然后我们得到下表：

index012345
A450-2-31
prefix0499745
prefix%50444340

其中 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。

代码：

    vector<int> cnt(K);
    for (auto n:prefix)
    {
        int num=(n%K+K)%K;
        cnt[num]++;
    }
    
    int ret=0;
    for (auto n:cnt)
        ret+=n*(n-1)/2;
    return ret;

这里我们用数组 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}。

这段代码和上面的那段连起来就是完整的解答了。

参考文献
[1] LeetCode 525.连续数组
https://www.cnblogs.com/geek1116/p/9389236.html

posted @ 2019-01-20 00:25  tomwillow  阅读(30)  评论(0)  编辑  收藏  举报