ABC167 D - Teleporter

atcoder.jp

Pythonです。精進しててえらいハマったので、なるべく分かりやすいソースにしてみたつもり。 ループしてなければ何事もなく終わるし、ループを検知したら、ループに戻る前にループサイクルを加味して、あと何移動すべきか決めたうえでループに飛び込んで終わる、、、ような感じ。 0-indexedに直すとそれはそれで混乱するし、悩ましい。

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

    checked= [0]*N

    Cur=1
    Next=0
    LoopDetect=False
    for i in range(K):
        if(checked[Cur-1]==0):
            checked[Cur-1]+=i
        Next=A[Cur-1]
        if(checked[Next-1]!=0):
            LoopDetect=True
            break
        Cur=Next


    if(LoopDetect):
        Cycle=i-checked[Next-1]+1
        amari=(K-i)%Cycle

        for i in range(amari):
            Next=A[Cur-1]
            Cur=Next

    print(Cur)

resolve()