Leetcode刷题·947和可被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

法一：暴力算法，遍历每个组合，算出每个组合的和是否能被K整除

class Solution(object):
    def subarraysDivByK(self, A, K):
        """
        :type A: List[int]
        :type K: int
        :rtype: int
        """
        res = 0
        for i in range(len(A)):
            sum = A[i]
            res += 1 if sum % K == 0 else 0
            for j in range(i+1,len(A)):
                sum += A[j]
                res += 1 if sum % K == 0 else 0
        return res

if __name__ == '__main__':
    s = Solution()
    re = s.subarraysDivByK([4,5,0,-2,-3,1],5)
    print(re)

法二：同余算法
同余定理：
若整数a和整数b除以正整数m的余数相等，则称a,b模m同余

我们可以知道，假设a<b,那么(b-a) mod m = 0
解决本问题就是看有没有子数组能被能被K整除，利用这个定理，那么我们可以设
count_mod_i = (A[0] + A[1] + A[2] +·······+A[i])%K,
count_mod_j = (A[0] + A[1] + A[2] +·······+A[i] + A[i+1] + ······ +A[j])%K,
这里 j>i,若count_mod_j - count_mod_i=0,那么从A[i]到A[j]这子数组必然能被K整除

class Solution(object):
    def subarraysDivByK(self, A, K):
        """
        :type A: List[int]
        :type K: int
        :rtype: int
        """
        dic = {}
        #初始化一个字典
        for i in range(K):
            dic[i] = 1 if i == 0 else 0
        for i in range(1,len(A)):
            A[i] += A[i-1]
        for a in A:
            dic[a%K] += 1
        return sum(v*(v-1) // 2 for v in dic.values())

为啥字典初始化时是dic{0:1}?
如果只有一个前缀和的余数是0那说明他也可以作为一个结果，这是与其它mod K不一样的，其它的需要两个才能算作一个结果，所以需要提前给0赋1，如果我们不提前给0赋1，那么当v是1的时候，也就是前缀和中只有一个数mod K是0的话，组合个数v*(v-1)//2结果是0，但我们是希望mod K为0是可以直接算作一个结果的，不然前缀和中只有两个mod K是0的时候才能算作一个，与题意不符。