2020-06-20

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)：周期を見抜く

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

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

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

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

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

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

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;
}