给定一个整数数组 nums
和一个整数 k
,返回其中元素之和可被 k
整除的非空 子数组 的数目。
子数组 是数组中 连续 的部分。
示例 1:
输入:nums = [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]
示例 2:
输入: nums = [5], k = 9 输出: 0
提示:
1 <= nums.length <= 3 * 10^4
-10^4 <= nums[i] <= 10^4
2 <= k <= 10^4
解法:前缀和 + 哈希表
通常,涉及连续子数组问题的时候,我们使用前缀和来解决。
我们令 P[i]=nums[0]+nums[1]+…+nums[i]。那么每个连续子数组的和 sum(i,j) 就可以写成 P[j]−P[i−1](其中 0<i<j)的形式。此时,判断子数组的和能否被 k 整除就等价于判断 (P[j]−P[i−1])modk==0,根据 同余定理,只要 P[j]modk==P[i−1]modk,就可以保证上面的等式成立。
因此我们可以考虑对数组进行遍历,在遍历同时统计答案。当我们遍历到第 i 个元素时,我们统计以 i 结尾的符合条件的子数组个数。我们可以维护一个以前缀和模 k 的值为键,出现次数为值的哈希表 record,在遍历的同时进行更新。这样在计算以 i 结尾的符合条件的子数组个数时,根据上面的分析,答案即为 [0..i−1] 中前缀和模 k 也为 P[i]modk 的位置个数,即 record[P[i]modk]。
最后的答案即为以每一个位置为数尾的符合条件的子数组个数之和。需要注意的一个边界条件是,我们需要对哈希表初始化,记录 record[0]=1,这样就考虑了前缀和本身被 k 整除的情况。
注意:不同的语言负数取模的值不一定相同,有的语言为负数,对于这种情况需要特殊处理。
哈希表使用长度为 k 的数组实现。
我们利用前缀和数组来求解问题,对于子数组 nums[i:j] (包含下标 i, j),其区间和为 sum[j]−sum[i - 1] (其中 sum 为预处理得到的前缀和数组),
我们要判断的是 (sum[j]−sum[i - 1])%K是否等于 0。
根据 mod 运算的性质,我们知道 (sum[j]−sum[i - 1])%K=sum[j]%K−sum[i - 1]%K。
故若想 (sum[j]−sum[i - 1])%K=0 ,则必有 sum[j]%K=sum[i - 1]%K。
所有满足 nums[i:j] 中元素之和可以被 K 整除的开始下标 i,必有 sum[j]%K=sum[i - 1]%K。我们以 sum[i - 1]%K 作为键值统计其出现的频率,从而对于每个下标 j 我们可以立即获得能和它组成满足要求的子数组的开始下标 i 的数量。
生成前缀和数组的过程中,将 key=sum[j]%k 出现的频率加入结果数组, 同时将 key=sum[j]%k 的出现频率加 1。
由于数组中有可能出现负数,我们需要将其加 K 从而使其 %K 之后的值为正数。
Java版:
class Solution {
public int subarraysDivByK(int[] nums, int k) {
int ans = 0;
int[] record = new int[k];
record[0] = 1;
int presum = 0;
for (int num: nums) {
presum += num;
int key = (presum % k + k) % k;
ans += record[key];
record[key]++;
}
return ans;
}
}
Python3版:
class Solution:
def subarraysDivByK(self, nums: List[int], k: int) -> int:
ans = 0
record = [0] * k
record[0] = 1
presum = 0
for num in nums:
presum += num
presum = (presum % k + k) % k
ans += record[presum]
record[presum] += 1
return ans
复杂度分析
- 时间复杂度:O(n),其中 n 是数组 nums 的长度。我们只需要从前往后遍历一次数组,在遍历数组的过程中,维护哈希表的各个操作均为 O(1),因此总时间复杂度为 O(n)。
- 空间复杂度:O(k),即 record数组 需要的空间。