水色でもわかる全方位木DP(rerooting)

参考にした記事

どれも1つだけでは理解できずにいろいろ読んでから手を動かしてみないと解らなかったので書く

何に使えるか

木の根から各subtreeについてある値を求めてその値をmergeする操作をsubtreeごとに再帰的に行う操作を木の全ての頂点を根として行いたい、みたいな場合。例として木の深さを求める場合など。 [4]によればmerge操作がモノイドであればよい(結合法則が成り立つ、単位元をもつ)とある(可換でもないといけない気がする？)

問題リストとしては [8] を参照。

基本的な考え方

[6]の図が一番解りやすかった。この赤と青の矢印を求めてしまえばあとは各頂点ごとに入ってくる矢印をmergeしてやればよい。赤の辺はある頂点を根としてDFSをしてやれば求められる。問題は青の辺をどう求めるかで、左右からの累積を使う方法と逆元を使う方法がある。

例として木の深さを各頂点から求める場合を考える。左の図がある木について頂点1からDFSした後の状態、各頂点で各subtreeのmaxを取って+1をしている。これを頂点2を根として行うと、1->2の辺は初めのDFSでは計算されていないため、これを求めておく必要がある。

f:id:isa_rentacs:20200405223615p:plain

逆向きの辺の求め方

左右からの累積を使う方法

上の木で頂点1から出る辺でまだ計算していないのは2,3,4への辺の値なので、これを求めることを考える。下の図で、一番左側は元の木の深さ1までの部分を取り出したもので、それ以外の図は1の子を親に持ってきた場合の木になっている。 f:id:isa_rentacs:20200405224038p:plain

これを見ると、1->2は3と4からの辺をmergeして1を加えたものとすればよく、1->3は2と4からの辺をmergeしたものに1を加えてやれば良さそう。これを一般化してやると、ある頂点から接続されている(注:子のみでないのに注意、親がいればそれも含まれる)頂点 $0 \dots N$ について、 $k$ 番目を親に持ってきたときの親に対する辺の値は $0\dots (k-1)$ 番目までの相手からの入力をmergeしたものと、 $(k+1)\dots N$ 番目までの相手からの入力をmergeしたものをさらにmergeしてやれば良さそう。

f:id:isa_rentacs:20200405224500p:plain

この左側の累積と右側の累積は相手の数Nに対してO(N)で求められる。一度目のDFSをした後にもう一度DFSをして各頂点についてこの操作を行うことで1回目のDFSと逆向きの辺の値がすべて埋まるのであとは各頂点を根としてmergeしてやればよい。

逆元を使う方法

累積を使う方法では、1つのsubtreeを取り除く(親に持ってくる)のにそれ以外のsubtreeの値をmergeすることで計算した。ここで、元のsubtreeすべてに対するmergeされた結果は1回目のDFSで計算されているので、1つのsubtreeをそこから取り除くことができればそれでもよい。 ABC160Fはこのケースで、例えば一番初めの木で頂点1を根としたときの場合の数は

$ans(1) = c_2 * c_3 * c_4 * {}_{n_2} C_{n_2} * {}_{n_2+n_3} C_{n_3} * {}_{n_2+n_3+n_4} C_{n_4}$

と書ける。 $c_i$ はiを根とするsubtreeの塗り方の数で、 $n_i$ はiを根とするsubtreeの頂点の数を表す。組み合わせの部分は $\frac{(n_2 + n_3 + n_4)!}{n_2!n_3!n_4!}$ と同様。ここで、例として頂点4を親に持ってきたとすると、計算したいのは頂点2と3を子としたときの場合の数で、それは

$c_2 * c_3 * {}_{n_2} C_{n_2} * {}_{n_2+n_3} C_{n_3}$

となる。ここで見比べてみると、頂点4を取り除くことで場合の数は

$\frac{1}{c_4 * {}_{n_2+n_3+n_4} C_{n_4}}$

倍になっていることが解る。 $n_2+n_3+n_4$ はすべての頂点数-1(根を除くため)なので、これは定数とみなせる。これが取り除く頂点に対応する逆元で、実際にはmodを取るのでmod_inverse(c_4, MOD) * mod_inverse(mod_comb(N, c_4, MOD)相当の計算をする必要がある。左右累積の場合と同様にこれをDFSで行うことで1回目のDFSと逆向きの辺すべてが求まる。

実装

ABC160F

左右累積 : Submission #11586064 - AtCoder Beginner Contest 160
逆元 : Submission #11577328 - AtCoder Beginner Contest 160

Isa@Diary

ソフトウェア開発やってます。プログラミングとか、US生活とかについて書きます。