D – Teleporter 解説 (AtCoder Beginner Contest 167)

AtCoder剰余, ダブリング, 周期性

問題へのリンク

問題概要

町が \(N\) 個ある。町 \(i\) から町 \(A_i\) に移動することを K 回繰り返す。
町 1 から始めた時、最終的にどの町にたどり着くか?

制約

  • \(2 \leq N \leq 2 \times 10^5\)
  • \(1 \leq A_i \leq N\)
  • \(1 \leq K \leq 10^{18}\)

考え方

単純にシミュレーションをしようとしても、\(K\) が大きいので計算が間に合いません。

簡単な例(例えば入力例2)などで試してみると様子がつかめます。

ある程度進んで、今までに到達したことのある町に到達したら、以降は周期的に同じ場所をグルグル回ることになります。

例えば、Aが

  • 6 5 2 5 3 2

の時、どのような移動の仕方になるかというと、

  • 1→6→(2→5→3)→(2→5→3)→(2→5→3)→(2→5→3)→…

というようになり、始めの「1→6」が終わったら、以降は「2→5→3」を繰り返すことになります。繰り返しの部分をシミュレーションする必要はなく、繰り返し部分の長さ(ここでは3)での余りを考えることで、必要な計算時間を大幅に減らすことができます。

もう一つの解法としては、ダブリングという概念を利用した方法もあります。

ダブリングは全体の要素数がN個のとき「K個先の要素を求めるのに \(O(K)\) かかる」ような状況を

  • 前処理:\(O(N \log K)\) 時間, \(O(N \log K)\) 空間
  • クエリ:\(O(\log K)\)

で行うことができるようにするアルゴリズムです。

解法1: 周期性と剰余を利用した方法

まずは同じ町が出現するまでシミュレーションをして、到達した順番を保存しておきます。少なくともN回移動すれば、今までに到達したことのある町に再び訪れることになります(鳩の巣原理)。よって、計算量は O(N) になります。

周期性のある場所に到達するまでを s、周期が r とすると、

  • K < s なら、K番目に到達する町を出力
  • s <= K なら、(K-s)%r + s 番目に到達する町を出力

C++ での実装例

問題文は 1-indexed ですが、C++ は 0-indexed なので、微妙にずらしてやる必要があります。

解法2: ダブリングを利用した方法

ダブリングを用いることでも解くことができます。

ここでは、

  • doubling[k][i] : 町 \(i\) から \(2^k\) 先の町はどこか?

という情報を前計算し、繰り返し二乗法のような2進数で見たときのビットに注目した計算をすることで、K個先の要素を簡単に取得することができます。(そもそも繰り返し二乗法がダブリングの1種とも言えます。)

  • 前処理:\(O(N \log K)\) 時間, \(O(N \log K)\) 空間
  • クエリ:\(O(\log K)\)

となるので、十分高速です。

C++ での実装例