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