LeetCode2845.统计趣味子数组的数目 -- 数学 | 模运算

1. 题目描述

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

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

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

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

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

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

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

2. 思路

首先，我们将原数组 nums 转化为一个 $01$ 数组：若 nums[i] % modulo == k，则转化为 $1$，反之为 $0$。

这样，我们就将问题转换为了，在一个 $01$ 数组 nums 中，有多少对下标对 $(l,r)$ 满足 (nums[l] + nums[l + 1] + ... + nums[r]) % modulo == k，我们可以用前缀和简化，也即满足 (s[r] - s[l]) % modulo == k。但枚举 $(l,r)$ 仍然需要 $O(n^2)$ 的时间。

注意到，本提保证 $0\le k<modulo$，所以 (s[r] - s[l]) % modulo == k 等价于：(s[r] - s[l]) % modulo == k % modulo。这意味着，s[r] - s[l] 与 k 关于模 $modulo$ 同余。根据模运算的性质，移项可得：(s[r] - k) % modulo == s[l] % modulo。

那么这下思路就简单了，遍历 nums，维护 s[l]，枚举 s[r]。

reference：灵茶山艾府

3. 模运算的性质

(a * b) % c = (a % c) * (b % c) % c;
(a + b) % c = ((a % c) + (b % c)) % c;
(a - b) % c = ((a % c) - (b % c)) % c;
(a / b) % c = (a * pow(b, c - 2)) % c;

reference: https://blog.csdn.net/Mikchy/article/details/81843306

4. 代码

// ref：灵茶山艾府
class Solution {
public:
    long long countInterestingSubarrays(vector<int>& nums, int modulo, int k) {
        long long res = 0;
        int n = nums.size();
        int s = 0; // 模拟前缀和
        // 注意到题目给出的 modulo 非常大(1e9)
        // 直接开 1e9 的空间肯定 OOM
        // 但是题目给定的 n 比较小(1e5)
        // 因此我们可以将 n+1 与 modulo 取一个小值开数组
        vector<int> f(min(modulo, (n + 1)));
        f[0] = 1;   // 预处理
        for(int i = 0; i < n; i ++ ) {
            s += (nums[i] % modulo == k);
            if(s >= k)
                res += f[(s - k) % modulo];
            f[s % modulo] ++ ;
        }
        return res;
    }
};