【ABC167 D】「Teleporter」を解く【Python3】

わがままな高橋王のために、テレポーターをK回使った後にどの町に到着しているかを求める問題です。

提出

n, k = map(int, input().split())
a = list(map(int, input().split()))

num = 1
li = [1]
flag = [True] * n
flag[0] = False

for i in range(k):
    num = a[num-1]
    if flag[num-1]:
        li.append(num)
        flag[num-1] = False
    else:
        break

d = li.index(num)
ans = (k-d)%(len(li)-d)+d

print(li[ans])

kが10¹⁸と大きいので、k回分の転送結果をその都度求めるには、時間が間に合いません。

町の数nが最大で「2×10⁵」であることに注目すると、（n>kであれば）k回転送している間に、一度訪れた町に再度転送され、そこからループすることが考えられます。

仮にループ内の要素の数が5個だった場合、5で割った余りを考えることで、k回分の転送を考える必要がなくなります。

提出のコードを具体的に説明していきます。

flagとして、すべての要素が「True」であるリストを作成しておきます。

町1からスタートなので、flag[0]は「False」にしておきます。

「num-1」番目に書かれている要素を「num」に入力していき、「num-1」番目のflagが「True」であれば、リスト「li」に付け加えて、「false」にします。

すでにflagが「False」であれば、for文を終了させます。

この結果、「li」の中には、ループが始まる前に訪れた町が入っています。

入力例1であれば、[1, 3, 4]

入力例2であれば、[1, 6, 2, 5, 3]

ループは入力例1では「1」から、入力例2では「2」から始まります。