【AtCoder】ABC241E Putting Candies Python解説

はじめに

ABC 241 E 問題 Putting Candies を解くために考えたこと、ACできるPython3(PyPy3)コードを紹介します。
また、今回は試験的にスライド形式にしています。筆者の記事をよく見てくださる方は、これまでとどちらが見やすいかコメントくださるとうれしいです。

E.Putting Candies

問題ページ
難易度 ： 水色 1248

考察

$K$ 回愚直に操作を追うことはできませんが、ループ構造に注目すると高速に最終値を求めることができそうです。

最終合計値を求めるためには、ループ構造を含む 3つの部分の合計値を求める必要があることがわかりました。
操作回数 $i$ を入力して、$i$ 回目の操作までの合計値を求められるテーブルを作成しておけば、これらを $O(1)$ で求めることができそうです。
終点の位置など、ループにまつわる情報を求めたいので、ループ構造を検知する必要があります。

状態の遷移を素直に前から見ていくと、初めて加算値が重複した位置でループ構造の終点を検出できます。同時に、その値をそれまでに加算していた位置で始点を検出できます。
したがって、加算値 $( mod\ N)$ を入力してそれが何回目の操作かを求められるテーブルを作成しておけば、これらを $O(1)$ で求めることができそうです。

操作回数は $N$ 回までありえるので、これを満たすテーブルを用意します。範囲外参照を絶対に起こさない程十分大きなサイズで余裕を持たせてもよいと思います。

コード

pypy3
124 ms/ 2000 ms

N,K=[int(nk) for nk in input().split()]
A=[int(a) for a in input().split()]
count_acc=[0 for _ in range(N+1)]
mod_count=[N for _ in range(N)]
mod_count[0]=0
for i in range(1,N+1):
    count_acc[i]=count_acc[i-1]+A[count_acc[i-1]%N]
    # 加算値が既出ならループ検出
    if mod_count[count_acc[i]%N]<N:
        end=i
        start=mod_count[count_acc[i]%N]
        break

    mod_count[count_acc[i]%N]=i


if end>K:
    ans=count_
    ans=count_acc[K]

else:
    # 1周目
    ans=count_acc[end]
    rem=K-end
    # ループ回数
    rep=rem//(end-start)
    ans+=(count_acc[end]-count_acc[start])*rep
    # 最終周
    rem%=(end-start)
    ans+=count_acc[start+rem]-count_acc[start]

print(ans)

補足

• 類題
  問題ページ
  ネタバレ防止で、abc 175 ～ 180 C ～ E と記します。
  こちらの問題も記事にしているのでぜひご覧ください(230221時点、未掲載 すぐにあげます)

終わり

質問やご指摘、ご意見はこちらまで
Twitter : Waaa1471