Codeforces Beta #84 Div1C Lucky Tree

概要

$N$頂点で構成された無向木が与えられる。 無向木の各辺には"Lucky"かそうでないかの属性がついている。 頂点から頂点までのパスにLuckyな辺が一つ以上含まれている時、それをLuckyなパスと呼ぶ。 異なる頂点$i,j,k$について、$i$から$j$までのパスと$i$から$k$までのパスが共にLuckyである$(i,j,k)$の通り数を求めよ。

制約

$1\le N\le 10^5$

考察

iを全頂点について試す。この時、$i$からLuckyなパスでつながっている頂点の数を$A(i)$とすると、 $A(i)*(A(i)-1)$が$(i,j,k)$の通り数となる。
$A(i)$をどう求めればよいだろうか？ 一見辺$e(u→v)$について、$e$を通らずに$v$と連結なすべての頂点の集合についてDPすればよさそうである。
しかしもし頂点の出次数が非常に大きい場合(例えばウニグラフの場合)計算量は$O(N^2)$となり、間に合わない。 (virtual中はこれを$O(N)$と勘違いしてしまいTLEした。)
そこで、「全方位木DP」と呼ばれるテクニックを使おう。
全方位木DPとは、まず単純な木DP(各頂点を親とした部分木についてのDP)を行い、その計算結果をもとに今度は各頂点を根とした木全体についてのDPをすることである。
まず適当な頂点$par$を根にして有向な木にしよう。 次に、$F(i):=i$を親とした部分木について、$i$からLuckyなパスでつながっている頂点の個数を算出しよう。
これは各$i$の子$j$について、$i$と$j$をつなぐ辺がLuckyな場合$C(j)$、そうでない場合$F(j)$の総和を求めればよい。
(ただし、$C(x)$とは、頂点$x$を親とした部分木の頂点の個数である。)
これは単なる木DPである。
今度は、$G(i):=i$を親とした部分木に含まれない頂点について、$i$からLuckyなパスでつながっている頂点の個数を算出したい。
$i$が$par$の場合、当然$G(i)=0$である。
そうでない場合、$i$の親$j$について、 $i$と$j$をつなぐ辺がLuckyな場合$C(par)-C(i)$、 そうでない場合、$G(j)+F(j)-F(i)$である。
理由は下の図を参考にしていただきたい。

後は、$A(i)=F(i)+G(i)$なので、一番最初に考察した式によって答えを算出すればよい。

手書き図見にくいな(だったらつくるな)