[974][中等][前缀和][哈希] 和可被K整除的子数组

题目描述

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

解题思路

本题与[523][中等][前缀和][哈希] 连续的子数组和相似, 只是由判断是否存在可以被整除(等价于正整数倍), 变为计算这样的数组的数量.

仍然使用哈希表, 哈希表的key值仍为余数, 值为这个前缀和余数出现的次数. 因此边遍历数组计算前缀和余数, 边判断当前余数是否已经在之前出现过, 如果已经出现过, 之前出现了多少次, 就可以组成多少个可以被K整除的连续的子数组, 作为最终结果的贡献. 然后更新这个哈希表, 继续推进到下一个位置.

在初始化的时候, 将余数为0的计数, 先初始化为1.

class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        mapping = {0: 1}
        resi, ans = 0, 0
        for num in A:
            resi += num
            resi %= K
            if resi in mapping:
                ans += mapping[resi]
            mapping[resi] = mapping.get(resi, 0) + 1
        return ans