和可被 K 整除的子数组——前缀和 kcnt[0] = 1是dummy表示自身 容易出错

974. 和可被 K 整除的子数组

难度中等51收藏分享切换为英文关注反馈

给定一个整数数组 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. 1 <= A.length <= 30000
2. -10000 <= A[i] <= 10000
3. 2 <= K <= 10000

class Solution(object):
    def subarraysDivByK(self, A, K):
        """
        :type A: List[int]
        :type K: int
        :rtype: int
        """    
        s = 0
        kcnt = collections.defaultdict(int)
        kcnt[0] = 1
        for n in A:
            s += n
            kcnt[s%K] += 1
        return sum(k*(k-1)//2 for k in kcnt.values())

        

　　

题目求连续子数组的和能被K整除，连续子数组的和就可以表示为前缀和的差，

比如 sum(A[i : j + 1]) = s[j + 1] - s[i]，

如果两个数的差能被K整除，就说明这两个数 mod K得到的结果相同，

只要找有多少对 mod k 相同的数就可以组合一下得到结果，

举例 对于A= [4,5,0,-2,-3,1] K = 5，

s = [0, 4, 9, 9, 7, 4, 5] ，

kcnt = [2, 0, 1, 0, 4] 代表有s中有两个元素的余数都为0（即0和5），1个元素的余数为2（即7），四个元素的余数为4（即4994）

所以在保证余数相同的情况下，取出两个数都可以得到一组答案。对于这个例子答案就是 C22 + C12 + C42 = 1 + 0 + 6 = 7

class Solution(object):
    def subarraysDivByK(self, A, K):
        """
        :type A: List[int]
        :type K: int
        :rtype: int
        """
        s = [0 for i in range(len(A) + 1)] #s代表前缀和，即s[i]表示sum(A[:i])
        kcnt = [0 for i in range(K)] #kcnt[i]代表s中有多少个元素 mod K 为i
        for i in range(len(A)):
            s[i + 1] = s[i] + A[i]
        for item in s:
            kcnt[item % K] += 1
        #print s, kcnt

        return sum(x * (x - 1) // 2 for x in kcnt)