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

给定一个整数数组 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]

方法一：

class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        res = 0  # 结果变量
        cur_sum = 0  # 记录到当前位置的前缀和，一开始还没处理，初始化为0
        # 用default可以避免key不存在时报错，不喜欢用额外库的同学也可直接用原生dict的get方法，设置默认为0
        presum_count = collections.defaultdict(int)  # 这里设为int，如果key不存在直接设置当前传入的key对应的值为0
        presum_count[0] = 1  # 初始状态，前缀和的套路，初始前缀和为0的情况出现一次
        for i in range(len(A)):
            cur_sum += A[i]  # 计算到当前位置前缀和
            for key in presum_count:  # 遍历dict
                if (cur_sum - key) % K == 0:
                    # 若当前前缀减去某key的mod为0，即说明当前位置前的某个位置的前缀和与当前位置前缀和相差若干个K(即满足题意)
                    # 则直接加上这个key对应出现的次数，因为题目范围过大，所以这个dict中的key可能会有很多，在for遍历查找时会超时
                    res += presum_count[key]
            presum_count[cur_sum] += 1
        return res

方法二：

class Solution:
    def subarraysDivByK(self, A: 'List[int]', K: int) -> int:
        res = 0
        pre_mod = 0  # 存储当前位置的上一个位置的前缀和的余数加上当前位置的值对K的余数
        presum_count = collections.defaultdict(int)
        presum_count[0] = 1
        for i in range(len(A)):
            pre_mod = (pre_mod + A[i]) % K
            # 下面这行感觉最不好理解
            res += presum_count[pre_mod]  # 如果能在dict中找到相同的pre_mod，说明当前节点前的某个位置的前缀和到当前位置的前缀和间存在若干个K
            presum_count[pre_mod] += 1
        return res

参考链接