题目描述

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

解题思路

方法一

使用回溯法，穷举所有子数组，计算其元素总和，判断是否能被 K 整除。但是这个算法时间复杂度为 O(n^2)，因此不能通过本题。

方法二

利用前缀和的思想，计算 A 的前缀和数组 S。则任意子数组 A[i:j] 的和可以表示为 S[j] - S[i-1]，其中 [i, j] 表示子数组的下标范围。

那么问题转化为，对于每个 S[j]，计算有多少个 S[i-1] 满足 S[j]-S[i-1] 能被 K 整除。

我们可以使用哈希表来存储所有的 (S[i] mod K, count) 键值对，其中 S[i] mod K 表示区间 [0, i] 内前缀和的余数，count 表示该余数出现的次数。我们从左往右遍历一遍数组 A，计算出当前位置的前缀和的余数 mod K，然后更新哈希表。这时，我们只需要在哈希表中查找 S[j] mod K 的差值，即 (S[j] mod K - (S[i] mod K)) mod K，然后将该差值的出现次数累加到我们的答案中即可。

代码实现

class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        count = {0: 1}  # 存储余数的出现次数，初始为 0 余数的出现次数为 1
        res, s = 0, 0
        for a in A:
            s += a
            res += count.get(s % K, 0)
            count[s % K] = count.get(s % K, 0) + 1
        return res

复杂度分析

时间复杂度

我们只需遍历一遍数组，时间复杂度为 O(n)。

空间复杂度

哈希表中存储了所有余数出现的次数，因此空间复杂度为 O(K)。