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

来源：力扣

解题思路

对于子数组的求和问题，我们可以转换为前缀和来求解
假如P[i]=A[0]+A[1]+...+A[i]表示从0到i的前缀之和，那么当我们需要求A[i]到A[j]的子数组之和，就可以转化为P[j]-Pi-1。为了一般化都会直接在初始位设置一个0
对于子数组和的求余可以再次转化。假如(A-B)%K=0,那么A%K=B%k

代码

class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        # c = Counter(i % K for i in [0, *accumulate(A)])
        # return sum(c[k] * (c[k] - 1) // 2 for k in c)

        #逐一统计
        # Sum, cnt = 0, 0
        # d = {0: 1} #初始0出现1次，是为了计算前缀和本身就可以整除K的情况
        # for a in A:
        #     Sum += a
        #     m = Sum % K
        #     s = d.get(m, 0) #计算其前面有多少个前缀和取模与之相等
        #     cnt += s #累加以i为结尾的子数组和可整除K的个数
        #     d[m] = s+1 #这一步放在最后面是不能用自身减自身
        # return cnt

        #单一统计，统计不同的求模，多少个前缀和的取模相等，假设为n，取C(2,n)即可
        d = {0: 1}
        Sum = 0
        for a in A:
            Sum += a
            m = Sum%K
            d[m] = d.get(m, 0) + 1
        cnt = 0
        for x, cx in d.items():
            cnt += cx*(cx-1)//2
        return cnt