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

気合で式を立てた。