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

分析：

在开始分析之前，我们先假定一个数组变量pre_sum用来保存输入数组A的前n项和，例如pre_sum[0]表示A[0], pre_sum[1] 表示A[0]+A[1]…依次类推。那么接下来，我们就以A=[4, 5, 0, -2, -3, 1], K=5为例，可以得到如下表格：

index	0	1	2	3	4	5
A	4	5	0	-2	-3	1
pre_sum	4	9	9	7	4	5
pre_sum%K	4	4	4	3	4	0

可以看到pre_sum中一些元素对K求模得到了相同的值，我们要知道当两个数a, b除以同一个数c得到相同的余数时，那么a和b之差一定可以被c整除，例如：
5 % 3 = 2， 11 % 3 = 2， 11 - 5 = 6 , 6 % 3 = 0.

我们再回到我们的表中, pre_sum[0] % K = 4, pre_sum[1] % K = 4, 所以pre_sum[1]-pre_sum[0]即A[1]可以被K整除，同理pre_sum[2]-pre_sum[0]即sum([A[1], A[2]])可以被K整除，pre_sum[2]-pre_sum[1]即A[2]可以被K整除，这实际上就是一个排列组合问题，相当于从n个数中抽出两个数有多少种可能，即C(n, 2),在这个样例中，4一共出现4次，所以求得C(4, 2) = 6, 我们会发现这和标准答案不符啊，标准答案应该是7才对啊，我们不要忘了，pre_sum[5]%K=0的情况，即便0只出现了一次，它依然算作一种情况

明白这一点之后，接下来我们要做的就是找出pre_sum%K中每一个数出现的次数，然后构造一个排列组合

代码

class Solution(object):
    def subarraysDivByK(self, A, K):
        result = 0    # 用来记录最终的结果
        length = len(A)
        pre_sum = {}    # 利用字典来记录每一个余数出现的次数，其中余数为键，它出现的次数为值
        curr_sum = 0
        # 构造字典pre_sum
        for i in range(length):
            curr_sum += A[i]    # 每次循环curr_sum就等于数组A的前i项和
            key = curr_sum % K
            pre_sum[key] = pre_sum.get(key, 0) + 1

        for key in pre_sum.keys():
            value = pre_sum[key]
            if key == 0:      # 键值为0，即余数为0
                result += value
            if value > 1:    # 余数出现次数大于1
                n = (value * (value - 1)) / 2
                result += n

        return result