Codeforces darkshadows's blog 「DP on Trees Tutorial」Problem 4

Codeforces darkshadows's blog 「DP on Trees Tutorial」Problem 3の続編です。
codeforces.com
問題4.
「木Tが与えられ、頂点iにはCiのコストが設定されています。根の頂点からスタートし、未探索の頂点へとランダムに移動していきます。未探索の頂点がなくなったら移動をやめます。コストの合計の期待値が最小になるようにするにはどの頂点を根とすれば良いでしょうか。」

まず、根を頂点1としましょう。頂点Vからスタートして部分木Vの頂点へと移動するときのコストの合計の期待値をf(V)としたとき、Vの子の頂点をv1, v2, ... , vnとして、

f(V)=Cv+∑(i=1~n){ f(vi) }/n

と表すことが出来ます。
頂点Vを総当たりしていけば答えが求まりますが、もっと高速な方法を考えましょう。
根を頂点Vとしたときにf(V)を計算する場合、すでに求めているf(v1), f(v2), ... , f(vn)の他に、
「根を頂点1としたときの頂点Vの親」をparent(V)としたとき、V→parent(V)へと移動できることも考える必要があります。
f:id:unosss:20200406141436p:plain

頂点Vを根としたとき、parent(V)からスタートして部分木parent(V)の頂点に移動していくときのコストの合計の期待値をg(V)と定めます。すると、頂点Vを根としたときのf(V)は、

f(V)={ g(V)+∑(i=1~n){ f(vi) } }/(n+1)

と計算することができますね。
さて、あとはg(V)を計算しましょう。
下の画像を見てください。
f:id:unosss:20200406142325p:plain

頂点pについて、親はp', 子はv1, v2, ... , vnであるとします。
青い線は、子→親へと向かう経路
赤い線は、親→子へと向かう経路です。
すると、頂点viについて

g(vi)=Cp+{ g(p)+f(v1)+f(v2)+...+f(vi-1)+f(vi+1)+...+f(vn) }/n

と計算することができますね。

①頂点1を根としたときに、全ての頂点iについてf(i)を求める(再帰的に計算できる)
②    〃      、全ての頂点についてg(i)を求める(根である頂点1から順に葉に向かって計算できる)

①→②の順に計算すれば可能です(実装略)。