\(\def \set #1#2{\{ #1 \ \vert \ #2 \}}\)
解法
\(a\) の累積和を \(s\) とおく． \(a\) の空でない連続する区間を \([l,r)\) とおく． ans は，\(\ (l,r)\ \) の組であって
\(s[r] - s[r] % K = r-l\) and
\(0 \leq l < r \leq N\)
\(0 \leq l < r \leq N\)
を満たすものの個数． 基本的には mod \(K\) で考えたいが，1つ目の式の右辺は mod \(K\) をとっていない事に注意． よって，1つ目の式は，
\(r-l < K\) and \(s[r] - s[l] \equiv r-l\) mod \(K\)
と同値．両辺を \(r,l\) の式に分離すると， \(s[r] - r \equiv s[l]-l\) mod \(K\).
実装
\(r-l < K\) の部分を一旦無視して \(s[r] - r \equiv s[l]-l\) mod \(K\) だけ考える．これは
\( for(r,N) \)
\( ans\) += \(cnt[s[r]-r\) mod \(K]; \)
\( cnt[s[r]-r\) mod \(K]\)++;
の形で求まる．
次に \(r-l < K\) の部分を考える． \(r-l < K\)を満たさない \(l\) に対して cnt[s[l]-l mod K]--; を施せばよい． これは，累積和の区間の幅や，queueを使って実装できる． Queue には，消す候補の map の key \(s[l]-l\) をストックしておく． \(K\) 個以上溜まったら削除に移る．
使っている記号，マクロ等 "https://ecsmtlir.hatenablog.com/entry/2022/12/23/131925"
int main() {
ll n,k;
cin >> n >> k;
vll s(n+1); rep(i,n) cin >> s[i+1], s[i+1] += s[i];
{// quque
queue<ll> q;
ll ans = 0;
map<ll,ll> cnt;
rep(r,n+1){
while(q.size() >= k) cnt[q.front()]--, q.pop();
ll x = (s[r] - r + n*k) % k;
ans += cnt[x];
cnt[x]++;
q.push(x);
}
cout << ans << endl;
}
return 0;
}
// -------------------------------------
{ // 累積和の区間の幅を利用
map<ll,ll> cnt;
ll ans = 0;
rep(r,n+1){
ll l = r-k;
if(l >= 0) cnt[(s[l]-l+n*k)%k]--;
ll x = (s[r] - r + n*k) % k;
ans += cnt[x];
cnt[x]++;
}
cout << ans << endl;
}
