AtCoder ABC 133 E - Virus Tree 2 - あるエンジニアのAtCoder奮闘記

irisruneです。個人的に以前のABCの4完が最近のABCの5完に相当するように思えるのですが、誰か統計取ったりしているのでしょうか。

問題

500点問題の割にアルゴリズムも実装もかなり軽めで、発想系の問題だと思いました。

import sys
import collections as col
import numpy as np
MOD = 1000000007


def input():
    return sys.stdin.readline().rstrip()


def main():
    n, k = [int(e) for e in input().split()]
    adj = [[] for _ in range(n+1)]
    for _ in range(n-1):
        u, v = [int(e) for e in input().split()]
        adj[u].append(v)
        adj[v].append(u)

    ans = 1
    queue = col.deque([(1, 0, k)])
    while len(queue) > 0:
        u, p, k_ = queue.popleft()
        ans = (ans * k_) % MOD
        if p == 0:
            k_next = k - 1
        else:
            k_next = k - 2
        for v in adj[u]:
            if v == p:
                continue
            queue.append((v, u, k_next))
            k_next -= 1
    print(ans)


main()

500点問題ということやグラフ（正確には木）問題ということもあってかなり難しく捉えがちですが実はかなり単純で、幅優先探索を行うことを考えると、

適当な点を始点として考えると、始点で選べる色は $K$ 通り。
始点から1番目に選んだ点で選べる色は $K-1$ 通り。
始点から2番目に選んだ点で選べる色は、
- 始点と隣接している場合は $K-2$ 通り。
- 始点から1番目に選んだ点と隣接している場合は $K-2$ 通り。
始点から3番目に選んだ点で選べる色は、
- 始点と隣接している場合は $K-3$ 通り。
- 始点から1番目に選んだ点と隣接している場合は $K-3$ 通り。
- 始点から2番目に選んだ点と隣接している場合は $K-2$ 通り。