けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 167 D - Teleporter (400 点)

 K 回操作した後の結果を求めよ (ただし  K がめちゃくちゃ大きい) という問題は、

  • 同じ処理の繰り返しとなっているところを見抜く
  • ダブリングする

というパターンが多いと思われる。

問題へのリンク

問題概要

 N 個の町  1, 2, \dots, N がある。町  i からは、町  A_{i} へテレポートできる。町 0 から出発して、 K 回テレポートしたときにいる町はどこになるか?

制約

  •  2 \le N \le 2 \times 10^{5}
  •  1 \le K \le 10^{18}

考えたこと

とにかく  K が大きいので、愚直にシミュレーションしたのでは間に合わない (愚直にシミュレーションしたら  O(K))。

こういうときの対処法として代表的なのが、

  1. きっと同じ処理の繰り返しになるはずなので、そのことを利用する
  2. ダブリングする

といったあたり。今回も両方でやってみたいと思う。

解法 (1):周期を見抜く

0-indexed で考えることにしよう。今回、頂点 0 から出発してたどっていくと、必ず

「一度来た頂点に戻ってくる瞬間」

があるはず。下図のグラフで考えてみる。

0 -> 2 -> 5 -> 11 -> 7 -> 1 -> 6 -> 12 -> 4 -> 8 -> 「5」

と進行していったときに、頂点 5 で「一度訪れた頂点」に初めて戻ってくることになる。

f:id:drken1215:20200620181002p:plain

そしてその後は、同じ動きをひたすらループすることになる。上図のグラフでは、

0 -> 2 -> 5 -> 11 -> 7 -> 1 -> 6 -> 12 -> 4 -> 8 -> 
5 -> 11 -> 7 -> 1 -> 6 -> 12 -> 4 -> 8 -> 
5 -> 11 -> 7 -> 1 -> 6 -> 12 -> 4 -> 8 -> 
5 -> 11 -> 7 -> 1 -> 6 -> 12 -> 4 -> 8 -> 
...

という風に繰り返すことになる。最初の方に「余計な何手か」がくっついてくるけど、それを除外すれば、完全に同じことの繰り返しになる。

よって、今回は、次のような解法で OK


  • まず頂点 0 から遷移を進めていく (ただし「一度来た頂点」を管理する)
  • 一度来た頂点に訪れたら、「最初の余計な分」が t 手分あるとして
    • K から t を引く
    • 「一度来た頂点のリスト」から、余計な分を取り除いてできる頂点のリストを a とする (a に含まれる頂点数を m とする)
    • このとき、a[ K % m ] が答え

具体的な実装は、頭がこんがらがりそうになるけど、落ち着いて実装することが大事そう。計算量は  O(K + N)

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

int main() {
    int N; 
    long long K;
    cin >> N >> K;
    vector<int> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i], --A[i];

    deque<int> a; // 繰り返しの部分
    vector<bool> seen(N, false); // 一度見たかどうか
    int cur = 0;
    while (true) {
        // 一度通った頂点を見つけたときの処理
        if (seen[cur]) {
            while (a[0] != cur) {
                // 最初の余計な数手分を除去する
                --K;
                a.pop_front();

                // 繰り返す前に K が限界になったらリターン
                if (K == 0) {
                    cout << a[0]+1 << endl;
                    return 0;
                }
            }
            break;
        }
        // 最初は愚直にシミュレーションしつつ、履歴をメモしていく
        a.push_back(cur);
        seen[cur] = true;
        cur = A[cur];
    }
    cout << a[K % a.size()]+1 << endl;
}

解法 (2):ダブリング

 K 回操作後の様子を求める方法として、非常に強力なダブリングとよばれる方法がある。この方法を使えば、今回の問題は思考停止で解くことができる。そのアイディアを簡単に書きたいと思う。

まず、

  • A[ v ] := v の 1 個先の町

という感じだったことを思い出そう。このことから、

  • B[ v ] = A[ A[ v ] ] という配列を用意すると、
  • B[ v ] := v の 2 個先の町

となることがわかる。さらに

  • C[ v ] = B[ B[ v ] ] という配列を用意すると、
  • C[ v ] := v の 4 個先の町 (2 個先の 2 個先なので)

となることがわかる。さらに

  • D[ v ] = C[ C[ v ] ] という配列を用意すると、
  • D[ v ] := v の 8 個先の町

となることがわかる。この考え方に基づいて、次のような配列を考えよう。


  • next[ 0 ][ v ] = A[ v ]
  • next[ d ][ v ] = next[ d - 1 ][ next[ d - 1 ][ v ] ]

つまり、next[ 0 ] は A に対応し、next[ 1 ] は上の B に対応し、next[ 2 ] は上の C に対応し、next[ 3 ] は上の D に対応する感じ。そして一般に、

  • next[ d ][ v ] := v の  2^{d} 個先の町

という風になる。もし  K が 2 の冪乗で、 K = 2^{m} という風に表せるならば、答えは単純に next[ m ][ 0 ] になる。より一般には、 K を二進法で表したときに


  •  d = 0, 1, 2, \dots に対して
  •  K d 桁目に 1 が立っているならば、
  • v = next[ d ][ v ] と更新する

という風にすれば OK。 d O(\log{K}) の分まで用意すれば OK。よって計算量は、 O(N\log{K}) となる。

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

int main() {
    int N; 
    long long K;
    cin >> N >> K;
    vector<vector<int>> next(60, vector<int>(N));
    for (int v = 0; v < N; ++v) cin >> next[0][v], --next[0][v];
    for (int d = 0; d + 1 < 60; ++d) {
        for (int v = 0; v < N; ++v) {
            next[d+1][v] = next[d][next[d][v]];
        }
    }
    int v = 0;
    for (int d = 0; d < 60; ++d) {
        if (K & (1LL<<d)) v = next[d][v];
    }
    cout << v + 1 << endl;
}