エイシングプログラミングコンテスト2021 E - Count Descendants 復習

あらすじ

E - Count Descendants

本番で解けなかったけどオイラーツアーとマージテクについての理解が深まったので良かった

問題概要

N頂点の木が与えられる。
Q個のクエリに答える。
クエリ: 「根(頂点1)からの距離がDかつ根への最短経路に頂点Uが含まれるような頂点はいくつあるか。」

この記事では、詳細な解説は公式解説に任せ、コンテスト中及び解説を見た後に考えたことの覚書を書いておきます。
自分以外に伝わりづらい表現、正確でない表現があらわれるかもしれませんがご了承ください。

コンテスト中に行った考察(TLE)

ひとまずDFSで全頂点の深さを求める。また、全頂点を深さごとにグルーピングする。
(例: 深さ2の頂点 => [2, 4, 8])

この状態で各クエリに以下の手順で答える

前処理で得た配列から深さDの候補頂点を全部取得する。
それぞれの候補頂点に対してLCA(候補, U) = U であればクエリの条件を満たすので ans += 1

この方法では計算量は O(QNlogN) になり、TLE
問題点は「毎回毎回すべての候補頂点をチェックする」というところにある。
たとえば根が中心になっているウニグラフでD=1のクエリが無数に来ると死ぬことになる。

200000
1 1 1 ... 1 1 1
200000
1 1
2 1
...
199999 1
200000 1

改善方法(公式解説の解法)

さて、公式解説の解法では「深さごとにグルーピングする」という方針までは同じですが、頂点番号ではなく対応する頂点のオイラーツアー行きがけ番号を入れています。

各クエリでは候補となる深さDの相当するの番号のうち、頂点Dの行きがけ番号 D_in と帰りがけ番号 D_out の間にあるものが条件を満たす。→二分探索が使えて O(QlogN)という話に落ち着きます。

今回のオイラーツアーの仕事として、作成した数列の順序がそのまま

Dより先に立ち寄る頂点 | Dの子孫の頂点 | Dより後に立ち寄る頂点

というグループ分けに対応しているので、あとはそのグループの境界を見つければいいことになります。境界の発見は二分探索の典型ですね。

というわけで解説を一読したときは天才か～と思いましたが、そのあと考えてみると他の問題にも適用できそうなテクニックが抽出出来て良かったです。
(オイラーツアーについてよく理解している方からは「そもそもそれがオイラーツアーだろ」と言われてしまいそうですが、自分が今まで理解していなかったので...
他の性質についてもコンテスト本番で驚かないようにこの機会に調べておきたいと思います)

マージテク解法

ところで、コンテスト中にあった別の考察として「前処理をすることによってどうにか各クエリO(1)で回答したい」というものがありました。

そのために各頂点について、「自分の子孫の中で、根からの距離がDの頂点がそれぞれいくつあるか」という情報を辞書などで保持しておけないかと考えました。

たとえば問題中サンプル1の頂点2について考えると、
「距離1: 1つ(2), 距離2: 3つ(4,5,7), 距離3: 1つ(6)」
という具合です。

f:id:socha77:20210523043056p:plain — 問題ページから引用しています

これは、葉となる頂点から順番に情報を吸い上げていくことで集約できます。
が、愚直にやるとO(N²)かかることがあります。

理解が滅茶苦茶直感的なものしかなくて申し訳ないのですが、「どうやら子孫に頂点 i が存在する」という情報は根にたどり着くまでに dist[i] 回伝播されることになります。
(なんというか、配列の値の更新がdist[i] 回行われるイメージ。
もちろん、深さdの頂点がk個のような形で上手く集約し、伝える情報の量が小さくなればその通りではない)

f:id:socha77:20210523044408p:plain — 伝わる？

というわけですべての頂点が一直線に並ぶような木であれば $\sum_{k=1}^n k$ でO(n²) 回の値の更新が必要になります。

ところで、このように値を集約していく際のテクニックとしてマージテクというものがあります。
これも自分の中だけで納得されている直感的な理解になってしまうのですが、

「二つのデータを合わせて新しいデータを作る」という操作を繰り返してデータを集約する
子データAには a 個の値、子データBには b 個の値がある(というより、値の更新に a ステップ、b ステップかかるというイメージ)
新しいインスタンスなどを用意してそこに両方の値を集約すると a+bステップかかるが、ポインタを共有するなどしてaにbを併合するようにすると、(1 + b) ステップで済む(1とかステップはどういう単位やねん)
いつも a > bとなるようにすると、最終的にO(NlogN) ステップで済む (1マージごとにサイズが 2b 以上になるからという証明の流れだったはず)

ちゃんと確認したい方は
"データ構造をマージする一般的なテク" とは？ - (iwi) ｛反省します - TopCoder部
などを参照してください。

これで前処理の計算量がO(NlogN) に抑えられます。うれしいね。

ところで...
実はこの方法は最終的に一つに集約することだけを考えているかつポインタをバリバリに共有しているので、根まで伝播が終わるころにはその子孫に置いておきたかった情報は壊れていることになります。