前缀和

leetcode974:和可被 K 整除的子数组

描述：给定一个整数数组 A，返回其中元素之和可被 K 整除的（连续、非空）子数组的数目

 

令P[i]=A[0]+A[1]+...+A[i]   每个连续子数组的和sum(i,j) 就可以写成 P[j] - P[i-1]（其中 0 < i < j0<i<j）的形式。此时，判断子数组的和能被 k 整除 就等价于判断 (P[j]−P[i−1])mod K == 0。

根据回余定理，只要 P[j] mod K == P[i-1]mod K，就可以保证上面的等式成立

答案就是 以每一个位置为数尾的符合条件的子数组个数之和

时间复杂度：从前往后遍历一次数组，再遍历数组的过程中，维护哈希表的各个操作均为O(1) ，总时间O(N)

原数组：A = [4,5,0,-2,-3,1], K = 5
前缀和：[0, 4, 4, 4, 2, 4, 0]
前缀哈希： Counter({4: 4, 0: 2, 2: 1})

 

/* 
let A = [4,5,0,-2,-3,1], K = 5
let map = {0: 1}
let preSum = 0, count = 0

循环A
i = 0,  preSum = (0 + 4) % 5 = 4, map中不存在4 => map[4] = 1
i = 1,  preSum = (4 + 5) % 5 = 4, map中存在4 => count += map[4] => count = 1,  map[4] = map[4] + 1 = 2
i = 2,  preSum = (4 + 0) % 5 = 4, map中存在4 => count += map[4] => count = 3,  map[4] = 3
i = 3,  preSum = (4 - 2) % 5 = 2, map中不存在2 => map[2] = 1
i = 4,  preSum = (2 - 3) % 5 = -1 小于0 => preSum = -1 + 5 = 4, map中存在4 => count += map[4] => count = 6,  map[4] = 4
i = 5,  preSum = (4 + 1) % 5 = 0, map中存在0 => count += map[0] => count = 7, map[0] = 2
*/

 

class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        #前缀和加取模运算加哈希;注意python是向下取整的
        hashmap = {0:1}
        presum = 0
        count = 0
        for i in range(len(A)):
            presum = (presum + A[i]) % K
            if presum in hashmap:
                count += hashmap[presum]
            if hashmap.get(presum):
                hashmap[presum] += 1
            else:
                hashmap[presum] = 1
            # print(hashmap)
        return count