計算したい内容は公式解説と同様です。
すなわち、適当な頂点を根として、各頂点 \(v\) とその子 \(u\) について、 \( |S_v| - \sum_{w \in S_v, c_w = c_v, f(P_{w, u}) = 2} |S_w| \) を計算したいです。ただし、各文字の定義は公式解説の通りです。

これは、 Auxiliary Tree を用いることで各色 \(c\) について高速に計算できます。詳細は実装例の main 部分を見てください。
計算量は、全体として \(O(N\log N)\) になります。

実装例