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⁴
-10⁴ <= nums[i] <= 10⁴
2 <= k <= 10⁴

解法

方法一：哈希表 + 前缀和

假设存在 \(i \leq j\)，使得 \(\textit{nums}[i,..j]\) 的和能被 \(k\) 整除，如果我们令 \(s_i\) 表示 \(\textit{nums}[0,..i]\) 的和，令 \(s_j\) 表示 \(\textit{nums}[0,..j]\) 的和，那么 \(s_j - s_i\) 能被 \(k\) 整除，即 \((s_j - s_i) \bmod k = 0\)，也即 \(s_j \bmod k = s_i \bmod k\)。因此，我们可以用哈希表统计前缀和模 \(k\) 的值的个数，从而快速判断是否存在满足条件的子数组。

我们用一个哈希表 \(\textit{cnt}\) 统计前缀和模 \(k\) 的值的个数，即 \(\textit{cnt}[i]\) 表示前缀和模 \(k\) 的值为 \(i\) 的个数。初始时 \(\textit{cnt}[0]=1\)。用变量 \(s\) 表示前缀和，初始时 \(s = 0\)。

接下来，从左到右遍历数组 \(\textit{nums}\)，对于遍历到的每个元素 \(x\)，我们计算 \(s = (s + x) \bmod k\)，然后更新答案 \(\textit{ans} = \textit{ans} + \textit{cnt}[s]\)，其中 \(\textit{cnt}[s]\) 表示前缀和模 \(k\) 的值为 \(s\) 的个数。最后我们将 \(\textit{cnt}[s]\) 的值加 \(1\)，继续遍历下一个元素。

最终，我们返回答案 \(\textit{ans}\)。

注意，由于 \(s\) 的值可能为负数，因此我们可以将 \(s\) 模 \(k\) 的结果加上 \(k\)，再对 \(k\) 取模，以确保 \(s\) 的值为非负数。

时间复杂度 \(O(n)\)，空间复杂度 \(O(n)\)。其中 \(n\) 为数组 \(\textit{nums}\) 的长度。

Python3JavaC++GoTypeScript

class Solution:
    def subarraysDivByK(self, nums: List[int], k: int) -> int:
        cnt = Counter({0: 1})
        ans = s = 0
        for x in nums:
            s = (s + x) % k
            ans += cnt[s]
            cnt[s] += 1
        return ans

class Solution {
    public int subarraysDivByK(int[] nums, int k) {
        Map<Integer, Integer> cnt = new HashMap<>();
        cnt.put(0, 1);
        int ans = 0, s = 0;
        for (int x : nums) {
            s = ((s + x) % k + k) % k;
            ans += cnt.getOrDefault(s, 0);
            cnt.merge(s, 1, Integer::sum);
        }
        return ans;
    }
}

class Solution {
public:
    int subarraysDivByK(vector<int>& nums, int k) {
        unordered_map<int, int> cnt{{0, 1}};
        int ans = 0, s = 0;
        for (int& x : nums) {
            s = ((s + x) % k + k) % k;
            ans += cnt[s]++;
        }
        return ans;
    }
};

func subarraysDivByK(nums []int, k int) (ans int) {
    cnt := map[int]int{0: 1}
    s := 0
    for _, x := range nums {
        s = ((s+x)%k + k) % k
        ans += cnt[s]
        cnt[s]++
    }
    return
}

function subarraysDivByK(nums: number[], k: number): number {
    const cnt: { [key: number]: number } = { 0: 1 };
    let s = 0;
    let ans = 0;
    for (const x of nums) {
        s = (((s + x) % k) + k) % k;
        ans += cnt[s] || 0;
        cnt[s] = (cnt[s] || 0) + 1;
    }
    return ans;
}

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

题目描述

解法

方法一：哈希表 + 前缀和

评论