LeetCode 974. 和可被 K 整除的子数组

974. 和可被 K 整除的子数组

给定一个整数数组 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数组 需要的空间。