974. 和可被 K 整除的子数组
题目描述
输入整数数组 nums 和整数 k。数组中可能有负数。
子数组指数组中一段连续且非空的元素。统计有多少个子数组的元素和能被 k 整除,返回这个数量。
示例 1
输入:nums = [4,5,0,-2,-3,1], k = 5 输出:7
解释:共 7 段,例如 [5]、[0]、[-2,-3]、[5,0] 以及整个数组等。
示例 2
输入:nums = [5], k = 9 输出:0
解释:唯一的子数组 [5] 和为 5,不能被 9 整除。
数据范围
1 <= nums.length <= 3 * 10^4-10^4 <= nums[i] <= 10^42 <= k <= 10^4
解题思路
子数组和 = 两个前缀和之差,可被 k 整除当且仅当两个前缀和模 k 相同。从左到右扫描,维护当前前缀模 k 的值 s,用哈希表 cnt 记录每个余数在前面出现过几次(空前缀余数 0 算一次)。每到一个新位置,先把 cnt[s] 加到答案(以这里结尾的合法子数组数),再把 cnt[s] 加一。
为什么正确
处理完 nums[0..i] 后,cnt[r] 是前缀 P[0..i+1] 中模 k 等于 r 的个数,ans 是右端点不超过 i 的合法子数组个数。
复杂度分析
时间 O(n),空间 O(k)。
易错点
负数取模:Python 的 % 结果在 [0, k) 内,但 C++/Java/Go/JS 中 -2 % 5 = -2,必须写成 ((s + x) % k + k) % k,否则余数 -2 和 3 会被当成不同的键。另外要先累加答案再增加计数,顺序颠倒会把空子数组也算进去。
参考代码
Python、Java、C++、Go、JavaScript、TypeScript 六种语言,逻辑一致。
class Solution:
def subarraysDivByK(self, nums: List[int], k: int) -> int:
cnt = {0: 1}
s = ans = 0
for x in nums:
s = (s + x) % k
ans += cnt.get(s, 0)
cnt[s] = cnt.get(s, 0) + 1
return ansclass Solution {
public int subarraysDivByK(int[] nums, int k) {
int[] cnt = new int[k];
cnt[0] = 1;
int s = 0, ans = 0;
for (int x : nums) {
s = ((s + x) % k + k) % k;
ans += cnt[s];
cnt[s]++;
}
return ans;
}
}class Solution {
public:
int subarraysDivByK(vector<int>& nums, int k) {
vector<int> cnt(k, 0);
cnt[0] = 1;
int s = 0, ans = 0;
for (int x : nums) {
s = ((s + x) % k + k) % k;
ans += cnt[s];
cnt[s]++;
}
return ans;
}
};func subarraysDivByK(nums []int, k int) int {
cnt := make([]int, k)
cnt[0] = 1
s, ans := 0, 0
for _, x := range nums {
s = ((s+x)%k + k) % k
ans += cnt[s]
cnt[s]++
}
return ans
}var subarraysDivByK = function(nums, k) {
const cnt = new Array(k).fill(0);
cnt[0] = 1;
let s = 0, ans = 0;
for (const x of nums) {
s = ((s + x) % k + k) % k;
ans += cnt[s];
cnt[s]++;
}
return ans;
};function subarraysDivByK(nums: number[], k: number): number {
const cnt: number[] = new Array(k).fill(0);
cnt[0] = 1;
let s = 0, ans = 0;
for (const x of nums) {
s = ((s + x) % k + k) % k;
ans += cnt[s];
cnt[s]++;
}
return ans;
}