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

解题思路

复习下，什么是前缀和。

• 从数组第0项到当前项之和。
• presum[0]，表示数组A第0项到第0项的总和；
• presum[1]，表示数组A第0项到第1项的总和；
  ...
  于是有，

presum[i] = A[0] + A[1] + A[2] + ...A[i]

数组某项可以表示成相邻前缀和之差：

A[i] = presum[i] - presum[i-1]

题意转化

求元素和能被K整除的子数组数目，就是有几种i,j组合，使得第i~j项之和modK==0
(presum[i] - presum[j]) mod K == 0
等价于
presum[i] mod K == presum[j] mod K
题目就变成了统计前缀和有相同模的组合个数。特殊的对于modK==0的前缀和，需要提前补充一下hash_map[0]=1。

实现

int subarrayDivByK(vector& nums, int K){
    int presum = 0, ret = 0;
    unordered_map mod_count;
    mod_count[0] = 1;
    for(int i=0; i