AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC167 D - Teleporter

問題原文

atcoder.jp

解法

 K の制約がとても大きいため、実際に  K 回の移動を試すわけにはいかない。

全ての町について、必ずどこかに町に移動できるのでいずれどこかのループをぐるぐる回ることになる。

ループは  N 回以内に必ず現れるので、どこがループするかを計算しておく。
これがあれば移動結果を高速に求められる。

ループに入る前に  K 回の移動を終えてしまう場合に少し注意する。

類題というか、ほぼ同じ問題が旧ABCに出ている。 atcoder.jp

実装

stackに積んでおいて、ループ始点が出るまで pop() とすると楽かなぁ。。。と思った。

N, K = map(int, input().split())
A = list(map(lambda x: int(x) - 1, input().split()))

stack = [0]
visited = {0}

# どこからループが始まるかを探す
now = 0
while A[now] not in visited:
    now = A[now]
    stack.append(now)
    visited.add(now)
roop_start = A[now]

# 実際にどこがループしているのかを探す
roop = []
while stack[-1] != roop_start:
    roop.append(stack.pop())
roop.append(roop_start)
roop = roop[::-1]
roop_size = len(roop)

# ループに入るまで移動を試す。K回移動してしまっても終わり。
now = 0
while now != roop_start and K:
    now = A[now]
    K -= 1

# ループの情報を元に答えを計算する。
if K == 0:
    print(now + 1)
else:
    print(roop[K % roop_size] + 1)

感想

diff700かぁ。。。みんな実装力あるなぁ。。。