974. Subarray Sums Divisible by K

https://leetcode.com/problems/subarray-sums-divisible-by-k/

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]

Note:

1. 1 <= A.length <= 30000

2. -10000 <= A[i] <= 10000

3. 2 <= K <= 10000

Thoughts

统计和是K的倍数的子数组个数。subarray sum => presum，DP或窗口。当presum出现相同余数时，说明中间子数组的和为K的倍数，a[i]+a[i+1]+...+a[j]=n1K+q; a[i]+a[i+1]+...+a[j]+...+a[n]=n2K+q => a[j+1]+...+a[n]=(n2−n1)K，因此用presum记录sum%K的频率，且presum添加默认的(0, 0)来处理前缀和刚好是K倍数的情况。

Code

class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        presum, s, res = collections.defaultdict(int), 0, 0
        presum[0] = 1
        for a in A:
            s += a
            s %= K
            res += presum[s]
            presum[s] += 1
        return res

class Solution {
public:
    int subarraysDivByK(vector<int>& A, int K) {
        unordered_map<int, int> prefix_freq;
        // The "null" before the first element. So when sum[0:i] % K == 0, the res is 1 not 0. 
        prefix_freq[0] = 1;
        int reminder = 0, res = 0;
        for (int a : A) {
            reminder = (a % K + reminder + K) % K;  
            res += prefix_freq[reminder]++;
        }

        return res;
    }
};