AtCoder Beginner Contest 133：E - Virus Tree 2

問題
解法
解答

問題

https://atcoder.jp/contests/abc133/tasks/abc133_e

解法

適当な頂点を根とする根付き木を考える．根から DFS をしていく．頂点 v にいるとき，その頂点の子の配色の通り数を考える．子は頂点 v の親と頂点 v の色それぞれと互いに異なる必要があるから，通り数は (K - 2)P(子の数)．ただし，v が根であるときは親が存在しなく，v 自身の色も決める必要があるから，通り数は K * (K - 1)P(子の数)．あとはこの通り数と各子の頂点の通り数を掛け合わせたものが頂点 v を根とする部分木全体の通り数となる．O(NlogK)．