TopCoder SRM 683 (Div. 2) Hard. SubtreesCounting
どことなく Codeforces を感じる。
解法
0 を根とする木を考えて木 DP をする。
- size[v] := 根が v であるような部分木のサイズの総和
- sub[v] := 根が v であるような部分木の個数
size[0]+...+size[n-1] が答えになる。
子が 3 つある頂点を根とする部分木は次のような形になっている。
そのため sub[i] は次のように計算できる。
$$ \begin{align} sub[i]&= 1+sub[c_0]+sub[c_1]+sub[c_2]+sub[c_0]sub[c_1] \\ &+ sub[c_0]sub[c_1]+sub[c_1]sub[c_2]+sub[c_0]sub[c_1]sub[c_2] \\ &=(1+sub[c_0])(1+sub[c_1])(1+sub[c_2]) \end{align} $$
size[i] の計算には、size[c0] が何回足されるを考えればいい。
上の図の子を左から c0, c1, c2 としたとき、c0 を含む部分木は 0,4,5,6 のパターンである。
- 0 のパターンのとき size[c0] は sub[c2] 回足される。
- 4 のパターンのとき size[c0] は sub[c1] 回足される。
- 5 のパターンのとき size[c0] は sub[c1]*sub[c2] 回足される。
- 6 のパターンのとき size[c0] は 1 回足される。
したがって size[c0] は 1+sub[c1]+sub[c2]+sub[c1]*sub[c2]=(1+sub[c1])(1+sub[c2]) 回足される。 少し式を変形すると
$$ \frac{(1+sub[c0])(1+sub[c1])(1+sub[c2])}{1+sub[c0]}=\frac{sub[i]}{1+sub[c0]} $$
になる。これを全ての子に対して考えてあげれば size[i] が求まって、
$$ size[i] = sub[i]\left(\frac{size[c0]}{1+sub[c0]}+\frac{size[c1]}{1+sub[c1]}+\frac{size[c2]}{1+sub[c2]}\right) + sub[i] $$
になる。一番最後に sub[i] を足しているのは i 自身が sub[i] 回足されるため。
後は dfs して順に求めていけばいい。
TopCoder SRM 683 (Div. 2) Hard. SubtreesCounting