ABC133 E - Virus Tree 2

問題原文

atcoder.jp

問題要旨

頂点からなる木の各頂点を、距離が２以下の頂点の色が異なるように 色の中から１色を選んで塗り分けるような方法は何通りあるか？

• 

解法

根つき木として、根から塗り分けていくことを考える。 図などを書けば、ある頂点とその直接の子はそれぞれ違う色で塗る必要があることがわかる。

また現在着目している頂点の親とも色が被ってはいけないので、その頂点と親が塗られているような頂点の、子の塗り分け方の通り数は、 となる。

ただし、親が存在しないような頂点（根）については、子の塗り分け方は となる。
また根は明らかに 色で塗り分けられる。

これは各頂点の深さを追っていくことで判別できる。

実装

MOD上のpermutaitonは逆元を使ってよしなにする。
)] が子の数ではないことに注意。

N, K = map(int, input().split())
T = [[] for i in range(N)]
MOD = 10 ** 9 + 7
for i in range(N - 1):
    a,  b = map(int, input().split())
    a, b = a - 1, b - 1
    T[a].append(b)
    T[b].append(a)

ans = K
stack = [[0, 0]]  # [頂点番号, 深さ]
visited = [False] * N
while stack:
    n, depth = stack.pop()
    visited[n] = True

    child = 0
    for to in T[n]:
        if visited[to]:
            continue

        child += 1
        stack.append([to, depth + 1])

    if depth == 0:
        ans *= nPr(K - 1, child)
    else:
        ans *= nPr(K - 2, child)
    ans %= MOD

print(ans % MOD)

感想

面倒な場合分けがいると思いきや、実は例外は少ないパターン、みたいな感じ？
どちらかというと実装で詰まりそうな問題。