AtCoder Beginner Contest 258 の E 問題。循環するリストの累積和のテクニック

話題とする問題はこちら。
知らなかったので、細かいところの実装コストが膨れて解けなかったので備忘のため。 atcoder.jp

本記事の主張内容はおおむね解説にある通りなので、当記事はお手軽に使える python スクリプトがあるよ、ということでお願いします。

E 問題を解くにあたって、以下のような部分問題を解くことになります。

$W_{0}, W_{1}, \dots, W_{N-1} \in \mathbb{N}$ および、 $X \in \mathbb{N}$ が与えられる。
各 $i =0, 1, \dots, N-1$ について、以下を満たす最小の $m$ を求めよ。
$X \leq \displaystyle{\sum_{i}^{i+m-1} W_{i \% N}}$

$m$ が問題に登場する「箱に入っているじゃがいもの数」となります。
これ単品であれば、ABC C 問題程度の難易度かなという印象です。

さて、思いつくのは累積和と二分探索による高速化ですが、悪い例（分岐が複雑）としてこんなスクリプトがあります。

import bisect

# 累積和の計算。s[i] は [0, i) の区間を足し合わせたものとする。
# 番兵の都合でリスト長は n + 1
s = [0] * (n + 1)
for k in range(n):
    s[k+1] = s[k] + w[k]
all_w = sum(w)

# i, i+1, ..., n-1, 0, ..., i-1 まで足した分は事前に足し上げておく
ans = (x // all_w) * n
# 事前に足し上げた分を取り除く
rem = x % all_w
if rem > all_w - s[i]:
    # n-1 まで足しても足りないので、とりあえず n-1 まで足す。
    ans += n - i
    rem -= all_w - s[i]
    # 0 からどこまで足せばよいか二分探索。
    j = bisect.bisect_left(s, rem, lo=0, hi=n+1)
    ans += j
else:
    # i からどこまで足せばよいか二分探索
    j = bisect.bisect_left(s, rem + s[i], lo=0, hi=n+1)
    ans += j - i
print(ans)

個人的には、「n-1 まで足しても足りないので、とりあえず n-1 まで足す」付近の ans や rem を更新するあたりで「１足すべき？　足さないで OK？」など細かいことを考えないといけなくなり、時間がかかってしまいました。

解説の提案する手法はこちらです。

import bisect

# 累積和の計算。s[i] は [0, i) の区間を足し合わせたものとする。
# リスト長を 2n にしておくのが肝。
s = [0] * (2 * n)
for k in range(2*n):
    s[k+1] = s[k] + w[k % n]
all_w = sum(w)

# i, i+1, ..., n-1, 0, ..., i-1 まで足した分は事前に足し上げておく
ans = (x // all_w) * n
# 事前に足し上げた分を取り除く
rem = x % all_w
# i からどこまで足せばよいか二分探索。
# リスト長を2倍にしているので、
# 添え字が n-1 から 0 に戻ってしまうのが１回以内ならば
# s の中で探しきることができる。
j = bisect.bisect_left(s, rem + s[i], lo=0, hi=n+1)
ans += j - i
print(ans)

コメントがごちゃごちゃしていてわかり辛いですが、分岐も減ってスクリプト量も減っており、何より二分探索の部分が $0 \leq i \leq j \leq N-1$ のときに書くスクリプトと、ほぼ同じスクリプトで OK になります。

累積和を、リスト長超えてぐるぐる足し合わせるような問題のときに、ぜひご活用ください。