E - Rem of Sum is Num Editorial
by
harurun4635
すべての \(A_i\) から \(1\) をひくことで、「要素の和を \(K\) で割った余りが \(0\) と等しくなるかつ長さが \(K\) 未満」と問題を言い換えることができ、見通しが良くなります。
長さが \(K\) 未満のみを管理するためには、\(j\) 番目の寄与を \(i = j + K\) 番目をみる直前に削除すればよいです。
実装例 (PyPy):
from collections import defaultdict
N, K = map(int, input().split())
a = [int(x) - 1 for x in input().split()]
s = [0] * (N + 1)
for i in range(N):
s[i+1] = (s[i] + a[i]) % K
cnt = defaultdict(int)
cnt[0] = 1
ans = 0
for i in range(1, N+1):
if K <= i: cnt[s[i-K]] -= 1
ans += cnt[s[i]]
cnt[s[i]] += 1
print(ans)
posted:
last update: