問題概要

あなたはアメーバの観察記録をつけました。最初 $1$ 匹のアメーバがおり、番号は $1$ です。

観察記録は時系列順に $N$ 個あり、 $i$ 番目の観察記録は

番号 $A_{i}$ のアメーバが分裂して消滅し、
新たに番号が $2i, 2i+1$ である 2 匹のアメーバが誕生した

というものです。このとき、アメーバ $A_{i}$ をアメーバ $2i,2i+1$ の親と呼びます。

各アメーバ $k=1, 2, \dots, 2N+1$ について何代親を遡るとアメーバ $1$ になるかを求めてください。

制約

$1 \le N \le 2 \times 10^{5}$

解法

Union-Find を自分の手で実装したことがある人なら、「頂点の親頂点」という概念に馴染みやすいのではないかと思う。今回の問題はつまり、

頂点 $A_{i}$ の子頂点が $2i, 2i+1$ である
頂点 $2i$ の親頂点は $A_{i}$ である
頂点 $2i+1$ の親頂点は $A_{i}$ である

というような根付き木を考えているのだ。根付き木の扱いに慣れている人なら、そう難しくはないはず。

根付き木の各頂点の深さを求める

今回の「何代親を遡るとアメーバ 1 になるか」は、要するに「根付き木の各頂点の深さを求めてください」ということだ。

たとえば、入力例 1 の $A = (1, 2)$ の表す根付き木は下図のようになる。頂点 1, 2, 3, 4, 5 の深さは、それぞれ 0, 1, 1, 2, 2 と求められる。

一般に、頂点 $v$ の深さを depth[v] と表すことにしよう。このとき、頂点 $v$ の親頂点が頂点 $p$ であるとき、

depth[v] = depth[p] + 1

という関係が成り立つ。よって、今回の問題では各 $i$ に対して、

depth[i*2] = depth[A[i]] + 1
depth[i*2+1] = depth[A[i]] + 1

という関係が成り立っている。この式を活用することで、各頂点の深さが求められる。

具体的には、次の実装例のように書ける。計算量は $O(N)$ となる。

注意点

注意しておきたいのは、depth[i*2] = depth[A[i]] + 1 や depth[i*2+1] = depth[A[i]] + 1 によって depth[i*2], depth[i*2+1] の値を求めるとき、depth[A[i]] の値がすでに求められているかどうかだ。

もし depth[A[i]] の値が求められていない状態だったら、depth[i*2] = depth[A[i]] + 1 や depth[i*2+1] = depth[A[i]] + 1 といった式を使うことが無意味になってしまうのだ。

しかし、それは大丈夫。制約条件に「観察記録は矛盾していない」とあるためだ。

実装例

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

int main() {
    int N;
    cin >> N;
    vector<int> depth(N*2+2, 0);
    
    for (int i = 1; i <= N; ++i) {
        int A;
        cin >> A;
        depth[i*2] = depth[A] + 1;
        depth[i*2+1] = depth[A] + 1;
    }
    
    for (int k = 1; k <= N*2+1; ++k) cout << depth[k] << endl;
}