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

Raphael Liu Lv10
  2023-09-13 00:47:26 2023-09-13 00:47:26 Created   2023-09-13 16:06:24 2023-09-13 16:06:24 Updated    

给你一个下标从 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 <= 105
  • 1 <= nums[i] <= 109
  • 1 <= modulo <= 109
  • 0 <= k < modulo

解决本题您需要掌握如下知识：

  1. 前缀和
  2. 哈希表
  3. 模运算

请看 视频讲解 第三题。

对于本题，由于需要统计 cnt，我们可以把满足 nums}[i]\bmod \textit{modulo} = k 的 nums}[i] 视作 1，不满足则视作 0。

如此转换后，算出 nums 的前缀和数组 s，那么题目中的 cnt}\bmod \textit{modulo} = k 等价于

(s[r+1]-s[l])\bmod \textit{modulo} = k

上式等价于（推导过程请看视频）

s[l]\bmod \textit{modulo} = (s[r+1]-k)\bmod \textit{modulo}

根据上式，我们可以一边枚举 r，一边用一个哈希表统计有多少个 s[r+1]\bmod \textit{modulo。这样可以快速知道有多少个 (s[r+1]-k)\bmod \textit{modulo，也就是 s[l]\bmod \textit{modulo 的个数，把个数加到答案中。

代码实现时，前缀和数组可以优化成一个变量 s。

[sol-Python3]
1
2
3
4
5
6
7
8
9
class Solution:
    def countInterestingSubarrays(self, nums: List[int], mod: int, k: int) -> int:
        cnt = Counter([0])  # 把 s[0]=0 算进去
        ans = s = 0
        for x in nums:
            s += x % mod == k
            ans += cnt[(s - k) % mod]  # Python 取模可以把负数自动转成非负数
            cnt[s % mod] += 1
        return ans
[sol-Java]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Solution {
    public long countInterestingSubarrays(List<Integer> nums, int mod, int k) {
        var cnt = new HashMap<Integer, Integer>();
        cnt.put(0, 1);  // 把 s[0]=0 算进去
        long ans = 0;
        int s = 0;
        for (int x : nums) {
            if (x % mod == k)
                s = (s + 1) % mod; // 这里取模，下面 cnt[s]++ 就不需要取模了
            ans += cnt.getOrDefault((s - k + mod) % mod, 0); // +mod 避免减法出现负数
            cnt.merge(s, 1, Integer::sum); // cnt[s]++
        }
        return ans;
    }

    // 数组版本，效率更高！
    // 因为 s 至多为 n
    public long countInterestingSubarrays(List<Integer> nums, int mod, int k) {
        int n = nums.size();
        var cnt = new int[n + 1];
        cnt[0] = 1;
        long ans = 0;
        int s = 0;
        for (int x : nums) {
            if (x % mod == k)
                s = (s + 1) % mod;
            int s2 = (s - k + mod) % mod;
            if (s2 <= n)
                ans += cnt[s2];
            cnt[s]++;
        }
        return ans;
    }
}
[sol-C++]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Solution {
public:
    long long countInterestingSubarrays(vector<int> &nums, int mod, int k) {
        unordered_map<int, int> cnt;
        cnt[0] = 1; // 把 s[0]=0 算进去
        long long ans = 0;
        int s = 0;
        for (int x: nums) {
            s += x % mod == k;
            ans += cnt[(s - k + mod) % mod]; // +mod 避免减法出现负数
            cnt[s % mod]++;
        }
        return ans;
    }
    
    // 数组版本，效率更高！
    // 因为 s 至多为 n
    long long countInterestingSubarrays(vector<int> &nums, int mod, int k) {
        int n = nums.size();
        vector<int> cnt(n + 1);
        cnt[0] = 1;
        long long ans = 0;
        int s = 0;
        for (int x: nums) {
            if (x % mod == k)
                s = (s + 1) % mod;
            int s2 = (s - k + mod) % mod;
            if (s2 <= n)
                ans += cnt[s2];
            cnt[s]++;
        }
        return ans;
    }
};
[sol-Go]
1
2
3
4
5
6
7
8
9
10
11
12
func countInterestingSubarrays(nums []int, mod, k int) (ans int64) {
	cnt := map[int]int{0: 1} // 把 s[0]=0 算进去
	s := 0
	for _, x := range nums {
		if x%mod == k {
			s = (s + 1) % mod // 这里取模，下面 cnt[s]++ 就不需要取模了
		}
		ans += int64(cnt[(s-k+mod)%mod]) // +mod 避免减法出现负数
		cnt[s]++
	}
	return
}
[sol-JavaScript]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
var countInterestingSubarrays = function (nums, mod, k) {
    const n = nums.length;
    var cnt = new Array(n + 1).fill(0);
    cnt[0] = 1;
    let ans = 0;
    let s = 0;
    for (const x of nums) {
        if (x % mod === k)
            s = (s + 1) % mod; // 这里取模，下面 cnt[s]++ 就不需要取模了
        const s2 = (s - k + mod) % mod; // +mod 避免减法出现负数
        if (s2 <= n)
            ans += cnt[s2];
        cnt[s]++;
    }
    return ans;
};

复杂度分析

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

相似题目（前缀和+哈希表）

推荐按照顺序完成。

 Comments
On this page
2845-统计趣味子数组的数目
  1. 相似题目（前缀和+哈希表）