Leetcode 974.和可被 K 整除的子数组

1 题目描述（Leetcode题目链接）

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

2 题解

  题目要求连续子数组，因此想到前缀和，那么中间一段 A [ i ] + ⋯ + A [ j ] A[i]+\cdots +A[j] A[i]+⋯+A[j]之和就可以通过前缀和相减求得。所以首先计算出前缀和，再遍历判断，这样的时间复杂度是 O ( n 2 ) O(n^2) O(n2)，可以做，不过超时了。

  那么就要思考一下如何优化，假如一段连续子数组为 A [ i + 1 , ⋯   , j ] A[i+1,\cdots,j] A[i+1,⋯,j]之和能够被 K K K整除，即
p r e f i x [ j ] − p r e f i x [ i ] = n k prefix[j]-prefix[i]=nk prefix[j]−prefix[i]=nk
其中 n n n为整数，简单移项并同时两边对 K K K取余我们就可以得到：
p r e f i x [ j ] % K = p r e f i x [ i ] % K prefix[j]\%K=prefix[i]\%K prefix[j]%K=prefix[i]%K
也就是两个前缀和取余 K K K相等。因此我们可以用哈希表来存储具有相同余数的前缀和了。

class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        res = 0
        prefix = 0
        dic = collections.defaultdict(int)
        for num in A:
            prefix += num
            i = prefix%K
            if i == 0:
                res += 1
            res += dic[i]
            dic[i] += 1
        return res

下面这样写也可以，方法是一样的。

class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        dic = collections.defaultdict(int)
        dic[0] = 1
        prefix = 0
        for num in A:
            prefix += num
            dic[prefix%K] += 1
        return sum([val*(val-1)//2 for val in dic.values()])