跳转至

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

来源第 361 场周赛 Q3难度中等分数2073

题目描述

给你一个下标从 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

解法

方法一：哈希表 + 前缀和

思考

子数组中满足 \(\textit{nums}[i]\bmod \textit{modulo}=k\) 的下标个数再模 \(\textit{modulo}\) 后等于 \(k\)。把满足条件的位置看成 \(1\)，问题化为前缀和 \(s\) 满足 \((s_r-s_{l-1})\bmod \textit{modulo}=k\)。哈希表统计前缀和模 \(\textit{modulo}\) 的出现次数，即可在扫描时累加配对。

题目要求一个区间内满足 \(\textit{nums}[i] \bmod \textit{modulo} = k\) 的索引 \(i\) 的数量，我们可以将数组 \(\textit{nums}\) 转换为一个 \(0-1\) 数组 \(\textit{arr}\)，其中 \(\textit{arr}[i] = 1\) 表示 \(\textit{nums}[i] \bmod \textit{modulo} = k\)，否则 \(\textit{arr}[i] = 0\)。

那么对于一个区间 \([l, r]\)，我们可以通过前缀和数组 \(s\) 来计算 \(\textit{arr}[l..r]\) 中 \(1\) 的数量，即 \(s[r] - s[l - 1]\)，其中 \(s[0] = 0\)。

我们用哈希表 \(\textit{cnt}\) 记录前缀和 \(s \bmod \textit{modulo}\) 出现的次数，初始时 \(\textit{cnt}[0]=1\)。

接下来，我们遍历数组 \(\textit{arr}\)，计算前缀和 \(s\)，将 \((s-k) \bmod \textit{modulo}\) 出现的次数累加到答案中，然后将 \(s \bmod \textit{modulo}\) 出现的次数加 \(1\)。

遍历结束后，返回答案即可。

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
class Solution:
    def countInterestingSubarrays(self, nums: List[int], modulo: int, k: int) -> int:
        arr = [int(x % modulo == k) for x in nums]
        cnt = Counter()
        cnt[0] = 1
        ans = s = 0
        for x in arr:
            s += x
            ans += cnt[(s - k) % modulo]
            cnt[s % modulo] += 1
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class Solution {
    public long countInterestingSubarrays(List<Integer> nums, int modulo, int k) {
        int n = nums.size();
        int[] arr = new int[n];
        for (int i = 0; i < n; ++i) {
            arr[i] = nums.get(i) % modulo == k ? 1 : 0;
        }
        Map<Integer, Integer> cnt = new HashMap<>();
        cnt.put(0, 1);
        long ans = 0;
        int s = 0;
        for (int x : arr) {
            s += x;
            ans += cnt.getOrDefault((s - k + modulo) % modulo, 0);
            cnt.merge(s % modulo, 1, Integer::sum);
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
public:
    long long countInterestingSubarrays(vector<int>& nums, int modulo, int k) {
        int n = nums.size();
        vector<int> arr(n);
        for (int i = 0; i < n; ++i) {
            arr[i] = int(nums[i] % modulo == k);
        }
        unordered_map<int, int> cnt;
        cnt[0] = 1;
        long long ans = 0;
        int s = 0;
        for (int x : arr) {
            s += x;
            ans += cnt[(s - k + modulo) % modulo];
            cnt[s % modulo]++;
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
func countInterestingSubarrays(nums []int, modulo int, k int) (ans int64) {
    arr := make([]int, len(nums))
    for i, x := range nums {
        if x%modulo == k {
            arr[i] = 1
        }
    }
    cnt := map[int]int{}
    cnt[0] = 1
    s := 0
    for _, x := range arr {
        s += x
        ans += int64(cnt[(s-k+modulo)%modulo])
        cnt[s%modulo]++
    }
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
function countInterestingSubarrays(nums: number[], modulo: number, k: number): number {
    const arr: number[] = [];
    for (const x of nums) {
        arr.push(x % modulo === k ? 1 : 0);
    }
    const cnt: Map<number, number> = new Map();
    cnt.set(0, 1);
    let ans = 0;
    let s = 0;
    for (const x of arr) {
        s += x;
        ans += cnt.get((s - k + modulo) % modulo) || 0;
        cnt.set(s % modulo, (cnt.get(s % modulo) || 0) + 1);
    }
    return ans;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
use std::collections::HashMap;

impl Solution {
    pub fn count_interesting_subarrays(nums: Vec<i32>, modulo: i32, k: i32) -> i64 {
        let mut arr: Vec<i32> = nums
            .iter()
            .map(|&x| if x % modulo == k { 1 } else { 0 })
            .collect();
        let mut cnt: HashMap<i32, i64> = HashMap::new();
        cnt.insert(0, 1);

        let mut ans: i64 = 0;
        let mut s: i32 = 0;

        for x in arr {
            s += x;
            let key = (s - k).rem_euclid(modulo);
            ans += *cnt.get(&key).unwrap_or(&0);
            *cnt.entry(s % modulo).or_insert(0) += 1;
        }

        ans
    }
}

评论