AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC146 E - Rem of Sum is Num

問題原文

E - Rem of Sum is Num

問題要旨

長さ N の整数列 A と整数 K が与えられる。Aの空でない連続する部分列であって、その和を K で割ったあまりが部分列の要素数と等しくなるようなものの個数を求めよ。ただし、 2 つの部分列が列として同じでも、取り出す位置が異なるなら区別して扱う。

考えたこと

  • O(N2)が許されるなら累積和でなんとかなりそうだけど。。。
  • なんか見たことありそうな雰囲気は感じるけれど。。。
  • 累積和は使いそうなんだけどな
  • なんもわからん

難しかったこと

  • 区間の大きさで条件が変わるのがむずかしい
  • とっかかりがつかめない。累積和とった後なにやればいいんだろう。。。

考えたかったこと

  1. 区間の大きさが1増えると、満たすべき   mod K の値が1増える → あらかじめ A の全ての要素から1ずつ引いておけば、その変化を相殺できる
  2. 区間の和は高速に必要になるので、とりあえず累積和をとる。これを  S として、 S_0 は0とする。
  3. mod K を考えているので、 S の各要素を mod K したものにしておく。
  4. mod K の取りうる値は、 0, 1, ... ,K - 1 なので、 長さK以上の区間は考えなくて良い。
  5. あとは累積和の基本に立ち返り、  S[r  +  1\rbrack -  S[l\rbrack ==  0 となるような  l を、各 r について求めれば良い。これは、A - Zero-Sum RangesD - Candy Distributionと同じようにできる。

実装

from itertools import accumulate
from collections import defaultdict
N, K = map(int, input().split())
A = list(map(int, input().split()))


# あらかじめ数列から1を引いておく
A = [a - 1 for a in A]

# mod K の累積和を取る
S = [0] + list(accumulate(A))
S = [s % K for s in S]


ans = 0
counter = defaultdict(int)
for i, s in enumerate(S):
    ans += counter[s]  # 過去の出現数を参考に数え上げ

    counter[s] += 1  # 今見たところを追加
    if i - K + 1 >= 0:
        counter[S[i - K + 1]] -= 1  # 長さK-1からはみ出たところを取り除く


print(ans)