経理のAI仕訳帳

AtCoder公開更新

AtCoderの関数グラフをどう解く？周期検出とダブリングを実装で理解

読了 約19分 たびすけ
AtCoderの関数グラフを学ぶ図。周期検出とダブリングの考え方

次に読む記事

関連するテーマの記事を、先に確認できます。

問題文から手法を選ぶ

AtCoderで関数グラフを見つけたら、まず「各状態の次の状態が一つだけ決まるか」を確認します。次に、追う回数と求めるものを見て手法を分けます。開始点から少しだけ進めばよいなら一回ずつ追い、同じ状態の繰り返しを利用して累積値を大きな回数ぶんまとめて更新するなら周期検出、配列全体を大きな回数だけ更新するならダブリングを選びます。

問題文で見る条件 選ぶ方法 素朴な計算量 改善後の計算量・追加メモリ
一つの開始点から経路を追い、再訪までの人数や経路を求める 一回追跡＋訪問済み O(N) ABC228 BはO(N)時間・O(N)メモリ
状態数がNに収まり、Kが大きく、累積値も求める 周期検出 O(K) ABC241 EはO(N)時間・O(N)メモリ
配列全体のK回後を求め、Kが非常に大きい ダブリング O(NK) ABC367 EはK > 0でO(N log K)時間、K = 0でO(N)時間・O(N)追加メモリ

ここでいう関数グラフは、有限個の状態それぞれに対して、有効な次の状態が一つだけある状態遷移です。状態から出る矢印が二本以上ある問題や、次の状態が存在しない問題は、同じコードを当てはめる前に別の手法を検討します。

関数グラフは「次が一つ」の状態遷移

状態を頂点、次の状態を矢印で表すと、各頂点の出次数が1になります。開始点から矢印をたどると、最初は一度しか通らない尾を進み、その先で同じ状態へ戻る周期に入ります。次が常に同じなので、一度同じ状態に戻れば、その後の状態列も同じ順序で繰り返します。

見るもの 意味 使い道
現在の状態 次の状態を決める添字や値 配列参照で一手進める
周期の入口 最初に周期へ入った時刻・状態 入口までの値と周期内の値を分ける
周期長 同じ状態が再び現れるまでの間隔 大きなKを周期単位でまとめる

自分自身を指す自己ループは、同じ状態を直ちに再訪するため周期長1です。ただし、ABC228 Bでは制約 Ai ≠ i により自己ループは入力に現れません。一般の関数グラフを考えるときの境界として覚えておきます。

一回ずつの追跡が遅くなる理由

一つの状態から次の状態へ一回進むだけなら、処理は単純です。しかしK回をそのまま繰り返す方法はO(K)時間になります。Kが1012や1018なら、入力の状態数が小さくても回数がボトルネックです。

関数グラフでは、状態を記録すれば高々N個の異なる状態を調べたところで再訪します。累積値を持たないABC228 Bなら訪問済み配列で止められます。ABC241 Eのように累積値を求める場合は、周期の入口までの合計、周期一周の合計、残りの端数を分けて保存します。配列全体を更新するABC367 Eでは、一回ずつの更新がO(NK)になるため、遷移そのものを2のべき乗回分にまとめます。

コードを読むためのPythonの前提

この3問では、入力の添字が1始まりか0始まりか、数値が添字なのかデータなのかを見分けます。次のPython公式教材で、コードに使う文法を確認できます。

mapは各要素へ関数を適用するイテレータなので、必要なときにlist(...)へ変換します。添字だけを0始まりへ直し、問題のデータとして与えられた個数は意味を変えない、という区別が3問に共通します。

訪問済みで止める：ABC228 B Takahashi’s Secret

ABC228 Bは、N人の友達と、最初に秘密を知る友達Xが与えられます。友達iが秘密を伝える相手はAiです。すでに秘密を知っている友達へ戻った時点で新しい人数は増えないため、再訪する前までの人数を数えます。

制約は 2 <= N <= 105、1 <= X <= N、1 <= Ai <= N、Ai ≠ i です。入力は N X の行と A1 ... AN の行、出力は最終的に秘密を知る友達の人数です。問題文はABC228 Bの公式問題、手順の説明は公式解説で確認できます。

再訪までを小さく追う

公式サンプル1の入力は次のとおりです。

4 2
3 1 1 2

1始まりでは 2 → 1 → 3 → 1 と進みます。友達2、1、3を初めて訪れ、友達1を再訪したので出力は 3 です。配列へ入れる前にXとAの各値を1減らすと、0始まりの経路は 1 → 0 → 2 → 0 になります。

入力から出力までのPython参考コードは次のとおりです。

N, X = map(int, input().split())
X -= 1
A = [value - 1 for value in map(int, input().split())]

seen = [False] * N
answer = 0

while not seen[X]:
    seen[X] = True
    answer += 1
    X = A[X]

print(answer)

seen[X]を調べてから訪問済みにし、A[X]へ移ります。各状態を二度処理しないので、制約の範囲では計算量はO(N)、追加メモリはseenのO(N)です。公式サンプル2の出力は 7 になります。

状態の周期と累積値を分ける：ABC241 E Putting Candies

ABC241 Eでは、最初はS0 = 0とし、時刻tの皿の中の総数をStとします。1回ごとに、StをNで割った余りを添字にしてA[rt]個を追加します。K回後に出力するのは、状態に使った余りではなく、皿の中にあるアメの総数です。

制約は 2 <= N <= 2 × 105、1 <= K <= 1012、1 <= Ai <= 106 です。入力は N K の行と、0始まりの A0 ... AN-1 の行です。問題文はABC241 Eの公式問題、状態関数と周期の説明は公式解説で確認できます。

時刻tの総数を St、状態を rt = St % N と置くと、次の加算は A[rt]、次の状態は (rt + A[rt]) % N です。状態だけはN通りなので、同じ状態が再び現れます。最初に状態rsが時刻sと時刻tで一致したら、周期長は p = t - s です。

周期に入る前の合計を Ss、周期一周の加算を C とすると、Kが入口より後なら q = (K - s) // p 周期と rem = (K - s) % p 回の残りに分けられます。答えは Ss + q × C に、周期の先頭からrem回分を足した値です。最後の答えへ % N は適用しません。

入口と周期を手計算する

公式サンプル1は、N = 5、K = 3、A=(2,1,6,3,1) です。アメの総数は 0 → 2 → 8 → 11 と変わり、出力は 11 です。参照した状態は 0 → 2 → 3、加算した値は 2, 6, 3 でした。別の公式サンプル2では、Kが 1012 で、出力は 826617499998784056 です。

次の小さな入力は公式サンプルではなく、周期の計算を確かめる机上の例です。N = 13、K = 13、A=(2,1,1,4,1,1,1,8,1,1,1,1,1) とします。状態は 0 → 2 → 3 → 7 → 2 なので、入口は時刻s = 1、周期長はp = 3です。入口までに2個を追加し、周期内の加算 1 + 4 + 8 = 13 を、残り12回分として4周します。したがって S13 = 2 + 4 × 13 = 54 です。

同じ状態に戻ると、その後に選ばれる添字と加算列も同じになります。これが、入口前・周期・余りを分けてよい理由です。

入力から出力までのPython参考コードは次のとおりです。

N, K = map(int, input().split())
A = list(map(int, input().split()))

seen_at = [-1] * N
prefix = [0]
total = 0

while seen_at[total % N] == -1:
    state = total % N
    seen_at[state] = len(prefix) - 1
    total += A[state]
    prefix.append(total)

repeat_time = len(prefix) - 1
cycle_start = seen_at[total % N]
cycle_length = repeat_time - cycle_start

if K <= cycle_start:
    answer = prefix[K]
else:
    cycle_sum = prefix[repeat_time] - prefix[cycle_start]
    cycles, remainder = divmod(K - cycle_start, cycle_length)
    answer = (
        prefix[cycle_start]
        + cycles * cycle_sum
        + prefix[cycle_start + remainder]
        - prefix[cycle_start]
    )

print(answer)

seen_atには状態が最初に現れた時刻を、prefixにはその時刻までのアメの総数を保存します。再訪時のrepeat_timeとcycle_startから周期長を求め、周期和を整数回だけ掛けます。状態を見つけるループは高々N+1回なので、計算量はO(N)、追加メモリはO(N)です。公式問題のKは1以上であり、一般化してK = 0を考える場合の初期値はS0 = 0です。

写像を2のべき乗で合成する：ダブリング

配列全体をK回更新する問題では、1回の更新がO(N)なので、素朴な方法はO(NK)時間です。ABC367 EのようにKが大きいときは、同じ写像を2のべき乗回適用した結果を順に作ります。

Pj[i]を「操作を2j回行った後、位置iが操作前のどの添字を参照するか」と定義します。0始まりの写像をXとすると、P0 = X、Pj+1[i] = Pj[Pj[i]] です。古い写像を二度合成するだけで、1回分が2回分、2回分が4回分へ倍になります。

最終的な写像Qは0回操作の恒等写像から始めます。Kの最下位ビットが1ならそのレベルのPjをQへ合成し、Kを右へ1ビットずらして次のレベルへ進みます。Kのビットが1のレベルだけを使うので、合成する回数はO(log K)です。

K回後の配列を求める：ABC367 E Permute K times

ABC367 Eは、数列XとAに対して、Bi = A[Xi]でAを置き換える操作をK回行い、最後の配列を出力する問題です。各Xiは1以上N以下の一つの添字を指します。

制約は 1 <= N <= 2 × 105、0 <= K <= 1018、1 <= Xi <= N、1 <= Ai <= 2 × 105 です。問題文はABC367 Eの公式問題、写像の合成と二進展開は公式解説で確認できます。

K = 13を二進展開する

小さな机上例として、N = 5、K = 13、X=(2,3,4,5,1)、A=(10,20,30,40,50) とします。0始まりのP0は (1,2,3,4,0)、P2は (4,0,1,2,3)、P3は (3,4,0,1,2) です。

13 = 1 + 4 + 8 なので、恒等写像QへP0、P2、P3を順に合成します。最初の二つを合わせると5回分で恒等写像へ戻り、最後にP3を適用したQは (3,4,0,1,2) です。初期Aをこの添字で読むと、出力は (40,50,10,20,30) になります。

この例では、K & 1で現在のビットを調べ、K >>= 1で次のビットへ進みます。写像の合成とKの二進展開が別々の処理ではなく、同じループの中でつながっています。

入力のXだけを1減らして、0始まりのjumpとして扱います。Aの値は添字ではないため減らしません。K = 0なら操作を一度も行わないので、恒等写像Qをそのまま使い、初期Aを出力します。

入力から出力までのPython参考コードは次のとおりです。

N, K = map(int, input().split())
jump = [x - 1 for x in map(int, input().split())]
A = list(map(int, input().split()))

Q = list(range(N))

while K > 0:
    if K & 1:
        Q = [jump[Q[i]] for i in range(N)]
    jump = [jump[jump[i]] for i in range(N)]
    K >>= 1

answer = [A[Q[i]] for i in range(N)]
print(*answer)

Qは、Kの二進展開で選んだレベルの写像を順に合成したものです。したがってQ[i]は、位置iがK回後に参照する初期配列Aの添字になります。最後にA[Q[i]]を読むことで、操作をK回行った後の値を各位置へ並べられます。

実装では、各2j回分の配列を保存せず、現在のjumpだけを次の2倍の写像へ更新します。Qも必要なときだけ更新するので、過去の写像表を保存せず、jumpとQを一つずつ持ちます。K > 0なら計算量はO(N log K)時間、K = 0なら恒等写像を使うためO(N)時間です。追加メモリはいずれもO(N)です。

公式サンプル1は、入力 N = 7、K = 3、X=(5,2,6,3,1,4,6)、A=(1,2,3,5,7,9,11) に対して、出力 7 2 3 5 1 9 3 になります。公式サンプル2はK = 0で、操作前の 4 3 2 1 がそのまま出力されます。公式サンプル3のようにKが1018でも、同じ対数回のレベルを調べます。

そのまま使える条件と、切り替える境界

3問のコードは、入力の数値をすべて同じ意味で扱っているわけではありません。次の表で、添字とデータを分けて確認します。

問題 0始まりへ直すもの そのまま扱うもの
ABC228 B 開始点Xと、行き先を表すAの各要素 なし
ABC241 E なし（Aはアメの個数） Aの個数、最終的なtotal。total % Nは参照用の状態だけ
ABC367 E 行き先を表すXの各要素 Aの値。これは最終出力のデータ

有限個の全状態について有効な次の状態が一つあることが、共通の前提です。ABC241 Eではtotal % Nで状態を0からN-1へ閉じ込められますが、終端や範囲外へ遷移する問題では、停止条件や別のデータ構造を先に設計します。一つの状態から複数の行き先を選ぶ一般グラフ、木、重み付きDPにも、関数グラフのコードをそのまま使いません。DFS、BFS、DPなど、分岐やコストを表せる手法へ切り替えます。

自己ループは一般の関数グラフなら周期長1として記録できますが、ABC228 BではAi = iが禁止されています。ABC367 EのK = 0は恒等写像、ABC241 Eの公式制約はK >= 1です。問題文の境界をコードの最初に書くと、停止条件と答えの意味を取り違えにくくなります。

3問を解くときの確認順

順番 問題 確認すること 次へ進む目安
1 ABC228 B 1始まりを0始まりへ直し、seenを付ける順序と再訪停止を追う 小さい経路を紙に書き、答えが初回訪問数になる理由を説明する
2 ABC241 E total % Nを状態として、入口・周期長・周期一周の加算・余りを分ける 最終出力が状態の剰余ではなくアメの総数だと説明する
3 ABC367 E Pjの二倍合成、K & 1、K >>= 1、K = 0を確認する 配列全体のK回更新を、必要なビットだけの写像合成へ置き換える

次に関数グラフの問題を見たら、まず「状態はいくつか」「次の状態は一つか」「K回後に何を出すか」を書き出します。一回の経路、周期を含む累積値、配列全体の大きな更新のどれに当たるかが決まれば、3問のどの実装から始めるかも決まります。

関連記事

この記事と関連するテーマから、次に読める記事を紹介します。

カテゴリの記事をすべて見る →