题目描述

难度分:$1762$

输入$n(1 \leq n \leq 2 \times 10^5)$、$k(1 \leq k \leq 10^9)$和长为$n$的数组$a(1 \leq a[i] \leq 10^9)$。

输出$a$的非空连续子数组$b$的个数，满足$sum(b) \% k=len(b)$。其中$sum$表示求和，$len$表示长度。

输入样例$1$

5 4
1 4 2 3 5

输出样例$1$

4

输入样例$2$

8 4
4 2 4 2 4 2 4 2

输出样例$2$

7

输入样例$3$

10 7
14 15 92 65 35 89 79 32 38 46

输出样例$3$

8

算法

同余+前缀和

设$s[i]=\Sigma_{j=1}^{i}a[i] \% k$，我们枚举子数组的右端点$r$，则需要找到$l$满足$(s[r]-s[l]) \% k=r-l$，这样子数组$(l,r]$就是个符合条件的$b$数组。而$sum(b) \% k \lt k$，所以$r-l \leq k-1$，可以进一步得到$(s[r]-s[l]) \% k=(r-l) \% k$，也就是说$s[r]-s[l] \equiv r-l(\% k)$。

这样一来就可以初始化一个映射$mp$，$mp[val]$表示$(s[i]-i) \% k=val$的索引位置（为了方便滑窗时收缩左边界，用一个双端队列存索引列表），初始情况下$mp[0]$中有一个$0$。然后枚举子数组的右端点滑动窗口，在保证左端点和右端点的距离不会大于$k-1$的情况下，只要发现$mp$中存在$(s[i]-i) \% k$，就把$mp[(s[i]-i) \% k]$的大小累加到答案上。

复杂度分析

时间复杂度

遍历一遍数组$a$预处理出前缀和$s$数组，时间复杂度为$O(n)$。遍历一遍$s$通过滑动窗口求出答案，时间复杂度为$O(n)$。因此，整个算法的时间复杂度为$O(n)$。

空间复杂度

前缀和数组$s$的空间复杂度是$O(n)$。$mp$映射中最多需要存$O(n)$个索引，所以空间复杂度也为$O(n)$。因此，整个算法的额外空间复杂度为$O(n)$。

C++ 代码

#include <bits/stdc++.h>

using namespace std;
typedef long long LL;

const int N = 200010;
int n, k, a[N], s[N];

int main() {
    scanf("%d%d", &n, &k);
    for(int i = 1; i <= n; i++) {
        scanf("%d", &a[i]);
        s[i] = (s[i - 1] + a[i]) % k;
    }
    unordered_map<int, deque<int>> mp;
    long long ans = 0;
    mp[0].push_back(0);
    for(int i = 1; i <= n; i++) {
        int r = ((s[i] - i) % k + k) % k;
        while(mp.count(r) && i - mp[r].front() > k - 1) {
            mp[r].pop_front();
            if(mp[r].empty()) mp.erase(r);
        }
        if(mp.count(r)) {
            ans += mp[r].size();
        }
        mp[r].push_back(i);
    }
    printf("%lld\n", ans);
    return 0;
}