leetcode: Subarray Sums Divisible by K

Problem

Given an array A of integers, return the number of (contiguous, non-empty) subarrays that have a sum divisible by K.

Example 1:

Input: A = [4,5,0,-2,-3,1], K = 5
Output: 7
Explanation: There are 7 subarrays with a sum divisible by K = 5:
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]

Solution

使用prefix sums的方法。
用一个presum变量记录当前位置 i 的总和$A[0]+...+A[i]$，取$key=presum\%K$。
那么对于任意 i 和$(A[0]+...+A[i])\%K==key$，若存在 j<i 有$(A[0]+...+A[j])\%K==key$，则等价于$(A[j+1]+...+A[i])\%K==0$。所以用一个hashmap来记录$presum\%K=key$的个数。
特别的，$(A[0]+...+A[i])\%K==((A[0]+A[i-1])\%K+A[i])\%K$

class Solution:
    def subarraysDivByK(self, A: 'List[int]', K: 'int') -> 'int':
        res = 0
        hashmap = {0:1}
        presum = 0
        for a in A:
            presum = (presum+a)%K
            
            if presum in hashmap.keys():
                res += hashmap[presum]
                hashmap[presum] += 1
            else:
                hashmap[presum] = 1
        return res