974. Subarray Sums Divisible by K
给你一个 array nums 和 k，问你总共有多少个 subarray 的总和是 k 的倍数
Example 1:
Input: nums = [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]
Example 2:
Input: nums = [5], k = 9
Output: 0
思路:
1.通常这种问你有几个 subarray 的题目都没办法真的 O(n^2) iterate nums[j~i]
这题也是 看范围会超过一点
这时就可以先分开想 要怎么对每个 nums[i] 算出以他为结尾的 subarray 中
符合条件 (nums[j~i] 总和是 k 的倍数) 的有几个
2.要算 subarray 的总和也就是 sum(nums[j~i]) 可以用老招 prefix sum
也就是 sum(nums[j~i]) = prefix[i] - prefix[j-1]
这时候可以发现如果要让 sum 是 k 的倍数
那么 prefix[i] 和 prefix[j-1] 模 k 之后必须要相等 这样才能把多出来的扣掉
那对于 i 来说也不用去 iterate 每个 j < i 检查 prefix[i] - prefix[j-1] 了
直接看 prefix[0] ~ prefix[j-1] 有几个人模 k 后和他一样就好
3.所以我们只需要算一次 prefix sum 然后从左到右 iterate i 一次
对 n=0~k-1 用 count[n] 纪录目前有几个 prefix[i]%k == n
之后往下检查 prefix[i+1] 时直接查 count[prefix[i+1]%k]
就可以知道 prefix[0] ~ prefix[i] 中有几个人模 k 后跟他一样
也就是有几个人可以跟他配成总和等于 0 的 subarray
4.如果 prefix[i]%k == 0 等于可以直接拿 nums[0~i] 来用所以要多加一个
也可以看成是 prefix[i] - prefix[-1] i 和 -1 配
所以一开始的 count[0] 要是 1 (代表 prefix[-1] = 0)
Python code:
class Solution:
def subarraysDivByK(self, nums: List[int], k: int) -> int:
count = [1] + [0]*(k-1)
sum = res = 0
for num in nums:
sum += num
res += count[sum%k]
count[sum%k] += 1
return res