ABC 167 D - Teleporter
問題
$i=1,\ldots,n$ のそれぞれについて行き先 $A _ i$ が決まっているとき、$1$ を $k$ 回移動したときの行き先を求めよ。（$n \leq 2 \times 10 ^ 5, \ k \leq 10 ^ {18}, 1 \leq A _ i \leq n$）
解説
$i$ を $2 ^ j $ 回移動したときの行き先を $DP[i][j]$ とすると、$2 ^ j $ 回の移動は$2 ^ {j-1} $ 回の移動を $2$ 回行うと考えて、
$$
DP[i][j] = DP[DP[i][j-1]][j-1]
$$
が成り立つ。
$k$ 回の移動は $k$ を $2$ 進数展開して
$$
k = 2 ^ {j _ 0} + \cdots + 2 ^ {j _ r} \ \ (j _ 0 < \cdots < j _ r)
$$
と表してから、$j = j _ 0, \ldots,j _ r $ について $i$ を以下のように更新する。
$$
i \leftarrow DP[i][j]
$$
実装
インデックスの都合で、補助的に $A _ 0 = 0$ とした。
n, k = map(int, input().split()) *A, = map(int, input().split()) A = [0] + A bit = list(bin(k)[2:][::-1]) m = len(bit) DP = [[None for j in range(m)] for i in range(n + 1)] for i in range(n + 1): DP[i][0] = A[i] for j in range(1, m): for i in range(n + 1): DP[i][j] = DP[DP[i][j - 1]][j - 1] i = 1 for j, b in enumerate(bit): if b == '1': i = DP[i][j] print(i)
