974. Subarray Sums Divisible by K

Title

给定一个整数数组 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

Solve

暴力

暴力的方法当然是最好想也是最好实现的，可惜会超时：

def subarraysDivByK_violence(self, A: List[int], K: int) -> int:ans = 0for i in range(len(A)):for j in range(i + 1, len(A) + 1):if sum(A[i:j]) % K == 0:ans += 1return ans

前缀和

一看到子数组和，马上就得想到用“前缀和”。

先来看一下什么是前缀和：从数组第0项到当前项的总和。
如果用一个数组preSum表示：

preSum[0]：数组A 第 0 项到第 0 项的总和

preSum[1]：数组A 第 0 项到第 1 项的总和

preSum[2]：数组A 第 0 项到第 2 项的总和

……

于是有：preSum[i]=A[0]+A[1]+…+A[i]preSum[i]=A[0]+A[1]+…+A[i]preSum[i]=A[0]+A[1]+…+A[i]
所以有：A[i]=preSum[i]−preSum[i−1]A[i]=preSum[i]−preSum[i−1]A[i]=preSum[i]−preSum[i−1]
可推出：A[i]+…+A[j]=preSum[j]−preSum[i−1]A[i]+…+A[j]=preSum[j]−preSum[i−1]A[i]+…+A[j]=preSum[j]−preSum[i−1]
考虑边界情况，当i=0时，i-1=-1，让preSum[-1]=0，可得：A[0]+A[1]+…+A[j]=preSum[j]A[0]+A[1]+…+A[j]=preSum[j]A[0]+A[1]+…+A[j]=preSum[j]
这样可以让边界情况也能套用通式。

题目中说的子数组的元素之和，就是数组第i项到第j项的和。

求子元素之和能被K整除的子数组数目，等价于求有几种i，j组合，使得i~j项的和mod K == 0，即满足(preSum[j]−preSum[i−1])modK==0(preSum[j]−preSum[i−1])modK==0(preSum[j]−preSum[i−1])modK==0等价于pre[j]modK==pre[i−1]modKpre[j]modK==pre[i−1]modKpre[j]modK==pre[i−1]modK

preSum 数组的每一项，怎么求？
- 前一项的前缀和+当前项=当前项的前缀和
- 求出preSum数组项让它mod K，mod完看哪两项相等，统计次数
- 通式有i,j两个变量，找出所有相等的两项，需要两层循环
引入哈希表
- 不关心前缀和对应数据A的哪一项，即不关心具体位置
- 关系出现过哪些前缀和%K的结果和对应的次数
- 用一个变量preSum，保存每次求出的前缀和%K，存入哈希表
  - key:前缀和%K，数值作为key
  - value:这个结果出现了几次

流程梳理

预置边界情况：map提前放入0:1，代表前缀和mod K == 0已经出现过1次
遍历数据A的每一项，求当前项的前缀和mod K，存入map中
- 如果之前没有存过，则作为key存入，值为1
- 如果之前已经存过，则对应的值+1（说明存在之前求出的前缀和，它mod K==当前前缀和 mod K，将之前求出的前缀和 mod K出现的次数累加给ans）

复杂度分析

时间复杂度：O(n)
空间复杂度：O(K)

Code

def subarraysDivByK(self, A: List[int], K: int) -> int:ans, preSum, modKMap = 0, 0, {0: 1}for i in range(len(A)):preSum = (preSum + A[i]) % K# 因为负数取模还是负数，所以需要加 Kif preSum < 0:preSum += Kif modKMap.get(preSum, None):ans += modKMap[preSum]modKMap[preSum] += 1else:modKMap[preSum] = 1return ans