LeetCode 2845. 统计趣味子数组的数目

2845. 统计趣味子数组的数目

给你一个下标从 0 开始的整数数组 nums ，以及整数 modulo 和整数 k 。

请你找出并统计数组中 趣味子数组 的数目。

如果 子数组 nums[l..r] 满足下述条件，则称其为 趣味子数组 ：

• 在范围 [l, r] 内，设 cnt 为满足 nums[i] % modulo == k 的索引 i 的数量。并且 cnt % modulo == k 。

以整数形式表示并返回趣味子数组的数目。

注意：子数组是数组中的一个连续非空的元素序列。

示例 1：

输入：nums = [3,2,4], modulo = 2, k = 1
输出：3
解释：在这个示例中，趣味子数组分别是： 
子数组 nums[0..0] ，也就是 [3] 。 
- 在范围 [0, 0] 内，只存在 1 个下标 i = 0 满足 nums[i] % modulo == k 。
- 因此 cnt = 1 ，且 cnt % modulo == k 。
子数组 nums[0..1] ，也就是 [3,2] 。
- 在范围 [0, 1] 内，只存在 1 个下标 i = 0 满足 nums[i] % modulo == k 。
- 因此 cnt = 1 ，且 cnt % modulo == k 。
子数组 nums[0..2] ，也就是 [3,2,4] 。
- 在范围 [0, 2] 内，只存在 1 个下标 i = 0 满足 nums[i] % modulo == k 。
- 因此 cnt = 1 ，且 cnt % modulo == k 。
可以证明不存在其他趣味子数组。因此，答案为 3 。

示例 2：

输入：nums = [3,1,9,6], modulo = 3, k = 0
输出：2
解释：在这个示例中，趣味子数组分别是： 
子数组 nums[0..3] ，也就是 [3,1,9,6] 。
- 在范围 [0, 3] 内，只存在 3 个下标 i = 0, 2, 3 满足 nums[i] % modulo == k 。
- 因此 cnt = 3 ，且 cnt % modulo == k 。
子数组 nums[1..1] ，也就是 [1] 。
- 在范围 [1, 1] 内，不存在下标满足 nums[i] % modulo == k 。
- 因此 cnt = 0 ，且 cnt % modulo == k 。
可以证明不存在其他趣味子数组，因此答案为 2 。

提示：

• 1 <= nums.length <= 10^5
• 1 <= nums[i] <= 10^9
• 1 <= modulo <= 10^9
• 0 <= k < modulo

---

提示 1

The problem can be solved using prefix sums.

---

提示 2

Let count[i] be the number of indices where nums[i] % modulo == k among the first i indices.

---

提示 3

count[0] = 0 and count[i] = count[i - 1] + (nums[i - 1] % modulo == k ? 1 : 0) for i = 1, 2, ..., n.

---

提示 4

Now we want to calculate for each i = 1, 2, ..., n, how many indices j < i such that (count[i] - count[j]) % modulo == k.

---

提示 5

Rewriting (count[i] - count[j]) % modulo == k becomes count[j] = (count[i] + modulo - k) % modulo.

---

提示 6

Using a map data structure, for each i = 0, 1, 2, ..., n, we just sum up all map[(count[i] + modulo - k) % modulo] before increasing map[count[i] % modulo], and the total sum is the final answer.

解法1：前缀和 + 哈希表

提示 1
这个问题可以用前缀和来解决。
提示 2
设 count[i] 是前 i 个索引中，nums[i] % modulo == k 的索引个数。
提示 3
count[0] = 0 和 count[i] = count[i - 1] + (nums[i - 1] % modulo == k ? 1 : 0) for i = 1, 2, ..., n。
提示 4
现在，我们要计算每个 i = 1, 2, ..., n，有多少个指数 j < i，使得 (count[i] - count[j]) % modulo == k。
提示 5
将 (count[i] - count[j]) % modulo == k 改写为 count[j] = (count[i] + modulo - k) % modulo。
提示 6
使用 map 数据结构，对于每个 i = 0、1、2、......、n，我们只需在增加 map[count[i] % modulo] 之前将所有 map[(count[i] + modulo - k) % modulo] 相加，总和就是最终答案。

Java版：

class Solution {
    public long countInterestingSubarrays(List<Integer> nums, int modulo, int k) {
        long ans = 0;
        int precount = 0;
        Map<Integer, Integer> map = new HashMap<>();
        map.put(0, 1);
        for (int num: nums) {
            precount += num % modulo == k ? 1 : 0;
            precount %= modulo;
            int key = (precount - k + modulo) % modulo;
            ans += map.getOrDefault(key, 0);
            map.merge(precount, 1, Integer::sum);
        }
        return ans;
    }
}

Python3版：

class Solution:
    def countInterestingSubarrays(self, nums: List[int], modulo: int, k: int) -> int:
        ans = 0 
        precount = 0
        dic = {0: 1}
        for num in nums:
            precount += 1 if num % modulo == k else 0
            precount %= modulo 
            key = (precount - k + modulo) % modulo 
            ans += dic[key] if key in dic else 0

            if precount in dic:
                dic[precount] += 1
            else:
                dic[precount] = 1
        return ans

复杂度分析

• 时间复杂度：O(n)，其中 n 为 nums 的长度。
• 空间复杂度：O(n)。