AtCoder ABC 167 D - Teleporter (400 点) in Python

問題の概要

N個の町があり、それぞれの町には1からNまでの番号が一意に割り当てられている（重複なし）。各町iにはテレポーターがあり、町Aiに移動することができる。このとき、町1から出発してK回の移動を繰り返したときの場所はどこでしょうか。 問題のURL

制約

• 2 <= N <= 2*10^5
• 1 <= Ai <= N
• 1 <= K <= 10^18

考え方

• 日本語が分かりにくいので、コードを見た方が早いかもしれません

制約を見て分かるように、Kの値が非常に大きいです。したがって、K回の操作を純粋にシミュレーションすることはできません。

そこで「どこかで循環するんだろうな～」ということまでは比較的容易に思いつくと思います。実際、町の数はN個しかないので、N回の移動（最初は1にいる）を行えば必ず循環します（鳩の巣原理）。したがって、循環する部分（循環が始まる地点）を特定してあげれば、求めるK回の移動後における地点を求めることができます。循環が始まる地点は以下のように求めることができます：

1. 辿った地点を覚えておいて、かつて辿ったことのある地点に到達するまで移動を繰り返す
2. N回の移動を繰り返したときの位置の変遷を覚えておいて、N回の移動時における地点に訪れたのは何回の移動を行ったときかを調べる

次に、得られた循環する部分を用いてK回の移動後の地点を求める方法について考えます。まず、Kが循環が始まる地点よりも小さい場合はそのまま出力してあげて大丈夫です。一方、Kが循環が始まる地点に達している場合、循環が始まる地点まで移動して、そこから残りの移動回数を循環する地点の数で割った余りの分だけ移動してあげればよいことが分かります。

解答（Python）

• 入力は0-indexになるようにしておきます

解法1：辿った地点を覚えていく方法

N, K = map(int, input().split())
A = [0] + list(map(int, input().split()))

# order[i]: 町iに到達したときの総移動回数
order = [-1] * (N + 1)

seen = []
pos = 1

while order[pos] == -1:
    order[pos] = len(seen)
    seen.append(pos)
    pos = A[pos]

period = len(seen) - order[pos] # 循環の周期
m = order[pos] # 循環する地点

if K < m:
    print(seen[m])
else:
    K = m + (K - m) % period
    print(seen[K])

解法2：事前にN回の移動履歴を覚えておく方法

N, K = map(int, input().split())
A = [0] + list(map(int, input().split()))

seen = [1]
pos = 1

for i in range(N * 2): # 後の参照のために余分に回しておく
    pos = A[pos]
    seen.append(pos)

loop_end = N
loop_start = N - 1

# 周期を見つける
while seen[loop_start] != seen[loop_end]:
    loop_start -= 1

period = loop_end - loop_start # 循環の周期
m = N - period # 循環が始まる地点

if K < m:
    print(seen[K])
else:
    K = m + (K - m) % period
    print(seen[K])

余談

解法1において次のように実装してしまうとTLEになるかもしれません。

seen = [1]
pos = 1
while True:
    if A[pos] in seen:
        seen.append(A[pos])
        break
    seen.append(A[pos])
    pos = A[pos]

一般に、pythonにおけるelement in listは非常に遅いことが知られています。注意しましょう。