统计趣味子数组的数目

题目

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

思路

因为本题$nums$数组中的元素只与取模结果是否为$k$有关，而与其本身的值无关，因此只需要对$nums$的元素对$modulo$取模，判断结果是否为$k$即可，如果为$k$我们将其改为$1$，反之改为$0$，然后我们用前缀和得到$nums$前$i$个数中取模结果为$k$的数的个数数组（即前缀和数组，记为$y$数组）。假如$nums[j...i]$中取模结果为$k$的数的个数$cnt$满足$cnt$ % $modulo==k$，对于前缀和数组$y$即有$y[i]-y[j-1]$ % $modulo==k$，即有$y[j-1]=(y[i]-k)$ %$mod$，因此用哈希表维护$y[i]$ % $mod$的值的个数即可。

代码

class Solution {
public:
    long long countInterestingSubarrays(vector<int>& nums, int mod, int k) {
        int n=nums.size();
        vector<int>y(n+1);
        for(int i=0;i<n;i++){
            nums[i]=nums[i]%mod==k?1:0;
            y[i+1]=y[i]+nums[i];
        }
        map<int,long long>mp;
        mp[0]++;
        long long ans=0;
        for(int i=1;i<=n;i++){
            ans+=mp[((y[i]-k)%mod+mod)%mod];
            mp[y[i]%mod]++;
        }
        return ans;
    }
};