ABC202-E Count Descendants (Blue-)
問題概要
サイズ$ N $の木が与えられる。次のクエリを$ Q $個処理せよ。
$ U, D $という値が与えられる。頂点$ U $を取って、根から距離が$ D $の頂点の個数を答えよ。
制約: $ N, Q \in [2, 2 \times 10 ^ 5], U \in [1, N], D \in [0, N - 1] $
考察
ダブリングで行ける?
LCAのこともあったのでダブリング的手法でいけるか?と思った。つまり、
node[i][j][k] := 頂点iから距離が(2^j)の頂点のk番目
みたいに保持させておいて、クエリごとに$ 101 _ {(2)} = 5 $ならば、
- 根から4マス先の頂点をリストに入れる。
- そのリストの中からそれぞれ1マス先を入れる。
で決めてみる。
しかし、ダブリングテーブル構築自体は$ O(N \log N) $だが、どうしてもQuery1回あたり$ O(N) $ほど最悪かかってしまう。
オイラーツアー
木に対して○○のテクをするといえばオイラーツアー。典型です。
この記事にいろいろな応用がある。maspyさんのやつ良い。
簡単にまとめると、
cnt = 0; void dfs(int now, int bef){ cntをメモ cnt++; for(子ども) dfs(子ども, now); cntをまたメモ cnt++; }
このように探索の始めと終わりにメモをする。DFSの性質上、ある頂点の2回のメモの値は$ b, e $ならば、子どもの頂点にまつわるメモの値は$ [b, e] $にちょうど入る。
また、オイラーツアーに直接関係あるわけではないが、$ b $をメモするときの何かしらの値を$ X _ b $、$ e $をメモするときのその値を$ X _ e $とすれば、$ X _ e $と$ X _ b $の差分を取れば、自分の子供全てに対しての操作を取得できる。
今回はどうするの?
今回の場合、オイラーツアーを知っていれば一発。
深さごとに$ b $に当たる値をメモして、そしてその同じ深さの中でソートされてれば(DFSされてるのでそもそもソートされてるけど)、 指定頂点の$ [b, e] $の区間内の$ b $に当たる数が、その指定の深さでの該当頂点の子どもの頂点の数に当たる。
これは二分探索lower_bound, upper_bound
を使えば、全体で$ O(N \log N) $で解ける。
これはそもそも深さ的あり得ないものがあったら打ち切っていたけど無くても動きそう。
よりよい解法 クエリ先読み
さきほど、$ X _ e $と$ X _ b $を見れば、子供全てについて行った結果がわかると書いたが、今回は先読みしてクエリを各頂点にあらかじめ置いておくことによって、$ X _ e $が分かった時点で、メモしておいた$ X _ b $との差分を追うことによって、DFSをしながら無駄なくクエリに答えられる。
簡単にメモする、クエリを置くとか言っていたけど具体的な実装はコードみるとわかりやすい。
一言
ikefumyありがとう(定期)。課題解決通話ありがとう!