ABC167 D - Teleporter - ペンギンかっこいい

備忘録

問題

回答

import sys
import os
import math
import bisect
import collections
import itertools
import heapq
import re
import queue
from decimal import Decimal

# import fractions

sys.setrecursionlimit(10000000)

ii = lambda: int(sys.stdin.buffer.readline().rstrip())
il = lambda: list(map(int, sys.stdin.buffer.readline().split()))
fl = lambda: list(map(float, sys.stdin.buffer.readline().split()))
iln = lambda n: [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)]

iss = lambda: sys.stdin.buffer.readline().decode().rstrip()
sl = lambda: list(map(str, sys.stdin.buffer.readline().decode().split()))
isn = lambda n: [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)]

lcm = lambda x, y: (x * y) // math.gcd(x, y)
# lcm = lambda x, y: (x * y) // fractions.gcd(x, y)

MOD = 10 ** 9 + 7
MAX = float('inf')


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    N, K = il()
    tel = [0] + il()

    # 初期位置
    position = 1

    # 移動回数がN回(町の数)以下の場合
    if K <= N:
        for k in range(K):
            position = tel[position]
        print(position)
        exit()

    # 移動回数がN回(町の数)以上の場合
    for n in range(N):
        position = tel[position]

    # 既にN回移動済みのため、
    # 必要な移動回数(K)から減算
    K -= N
    # 今いる地点(position)に再到達するまでに
    # 必要な移動回数を探索
    base = position
    cnt = 1
    while base != tel[position]:
        position = tel[position]
        cnt += 1

    # Kからループ数の余りを求め、
    # 余りの回数だけ移動を行う
    for k in range(K % cnt + 1):
        position = tel[position]
    print(position)


if __name__ == '__main__':
    main()

考え方

移動回数(K)が町の数(N)以下の場合には、
最大でも2 * 10^5回程度の移動なので、実際に移動をシミュレートすることで回答を得られる。

問題は移動回数が最大の10^18回の場合。
素直にシミュレートすると明らかに実行時間制限に引っかかってしまう。
そのため、計算に工夫が必要になる。

1 ~ Nまでの町をある一定回数移動を行うと、
既に訪れたことのある町に移動する。
そのため、あるタイミングから移動する町はループすることがわかる。

# 例:町に設置されたテレポーターが 6 5 2 5 3 2 の場合、
# 町1からスタートすると以下のような移動になる
1 → 6 → 2 → 5 → 3 → 2 → 5 → 3 → 2 ...
# はじめは1 → 6 → 2と移動し、あとは5 → 3 → 2の繰り返し

町の移動が一定のタイミングからループするということは、
ループ中の移動は考慮の必要がなく、
ループが終了した後、何回移動するかに注目すればよい。
そのため、移動回数(K)からループ時に移動する回数を割ったときの余りだけ、
実際にシミュレートすることで回答を得られる。

注意点として、ループが始まるまでに必要な移動回数をKから引いておく必要がある。
(上記の例では5 → 3 → 2の繰り返しが行われるまでに1 → 6 → 2と移動しているため、
全体の必要な移動回数からループ前の移動数を引いておく必要がある）

町の数だけ移動のシミュレートを行えば、
あるタイミングから必ずループする。
そのため、実装ではN回の移動シミュレートを行ったあと、
必要な移動回数(K)からNを減算している。