DFS と Wavelet Matrix を組み合わせた解法

問題は、ある頂点 \(U_i\) かその子孫の中（部分木）で、深さが \(D_i\) の頂点の個数を求める問題と言い換えることが出来ます。

これを解くためにまず、与えられた木をDFSの行きがけ順に頂点番号を記録して、1次元配列にすることを考えます。

入力例1の木はDFSを行うとき、1, 2, 4, 6, … の順番に探索されるので、以下のように表現できます。

\( ​A = \{1, 2, 4, 6, 5, 7, 3\}\)

更にこれに対応して、それぞれの頂点の深さを記録した配列を考えます。以下のようになります。

\(D= \{0, 1, 2, 3, 2, 2, 1\}\)

頂点1の深さは0、頂点6の深さは3、の様に対応しています。

さて、例えば頂点2の子孫は \(\{2, 4, 6, 5, 7\}\) ですが、これは配列 \(A\) の連続部分列です。
この範囲に対応する \(D\) の連続部分列は \(\{1, 2, 3, 2, 2\}\) です。この連続部分列に含まれるある値の個数が高速に計算できれば良さそうです。

つまり、ある頂点の \(A\) での位置と、それ以下のノード数が分かれば、1次元の範囲クエリに置き換える事ができます。 ある頂点の \(A\) での位置とそれ以下のノード数は DFS をしている時に構築することが出来ます。これはオイラーツアーというテクニックとして知られています。

さて、配列\(A\)の連続範囲 \([L, R)\) に含まれる値 \(V\) の個数を高速に求められれば良いことが分かりましたが、これは Wavelet Matrix の range freq クエリを用いることで \(O(\log N)\) で求めることが出来ます。
Wavelet Matrix/Tree の説明

range_freq(A, L, R, x, y) クエリとは、配列 \(A\) の範囲 \([L,R)\) に含まれる \([x,y)\) を満たす値の出現数を求められるクエリです。 これで \([D_i, D_i+1)\) の値の出現数を求めることができます。

実際のコード例はこちら の通りです。
Wavelet Matrix のライブラリパートを除いて、約30行程度で求められることが分かります。