D – Teleporter 解説 (AtCoder Beginner Contest 167)

2020年5月11日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 なので、微妙にずらしてやる必要があります。

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

int main() {
    ll N, K;
    cin >> N >> K;
    vector<int> A(N);
    for (int i = 0; i < N; i++) {
        cin >> A.at(i);
        A[i]--;  // 0-indexed に変更
    }

    set<int> used;     // 探索済みの町を保存
    vector<int> path;  // 探索の順番を格納
    int now = 0;
    while (used.count(now) == 0) {  // 探索済みの町に到達するまで
        path.push_back(now);
        used.insert(now);
        now = A[now];
    }

    int s = find(path.begin(), path.end(), now) - path.begin();  // s: ループする最初の部分に到達するまで
    int r = (int)path.size() - s;                                // 周期

    if (K < s) {  // ループ部分に到達しない時
        cout << path[K] + 1 << endl;
        return 0;
    } else {  // ループ部分に到達する時
        cout << path[(K - s) % r + s] + 1 << endl;
        return 0;
    }
}

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

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

ここでは、

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

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

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

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

C++ での実装例

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

int main() {
    ll N, K;
    cin >> N >> K;
    vector<int> A(N);
    for (int i = 0; i < N; i++) {
        cin >> A.at(i);
        A[i]--;  // 0-indexed に変更
    }

    int logK = 1;
    while ((1LL << logK) <= K) logK++;
    // doubling[k][i] : i番目から 2^k 進んだ町
    vector<vector<int> > doubling(logK, vector<int>(N));
    for (int i = 0; i < N; i++) {
        doubling[0][i] = A[i];
    }

    // 前処理 doubling の計算
    for (int k = 0; k < logK - 1; k++) {
        for (int i = 0; i < N; i++) {
            doubling[k + 1][i] = doubling[k][doubling[k][i]];
        }
    }

    int now = 0;
    for (int k = 0; K > 0; k++) {
        if (K & 1) now = doubling[k][now];
        K = K >> 1;
    }
    cout << now + 1 << endl;
}

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