AtCoder Beginner Contest 133 E - Virus Tree 2 （dfs）

 

E - Virus Tree 2

---

Time Limit: 2 sec / Memory Limit: 1024 MB

Score : 500500 points

Problem Statement

You are given a tree with NN vertices and N−1N−1 edges. The vertices are numbered 11 to NN, and the ii-th edge connects Vertex aiai and bibi.

You have coloring materials of KK colors. For each vertex in the tree, you will choose one of the KK colors to paint it, so that the following condition is satisfied:

• If the distance between two different vertices xx and yy is less than or equal to two, xx and yy have different colors.

How many ways are there to paint the tree? Find the count modulo 1 000 000 0071 000 000 007.

 

What is tree?

 

 

What is distance?

 

Constraints

• 1≤N,K≤1051≤N,K≤105
• 1≤ai,bi≤N1≤ai,bi≤N
• The given graph is a tree.

---

Input

Input is given from Standard Input in the following format:

NN KK
a1a1 b1b1
a2a2 b2b2
..
..
..
aN−1aN−1 bN−1bN−1

Output

Print the number of ways to paint the tree, modulo 1 000 000 0071 000 000 007.

---

Sample Input 1 Copy

Copy

4 3
1 2
2 3
3 4

Sample Output 1 Copy

Copy

6

There are six ways to paint the tree.

---

Sample Input 2 Copy

Copy

5 4
1 2
1 3
1 4
4 5

Sample Output 2 Copy

Copy

48

---

Sample Input 3 Copy

Copy

16 22
12 1
3 1
4 16
7 12
6 2
2 15
5 16
14 16
10 11
3 10
3 13
8 6
16 8
9 12
4 3

Sample Output 3 Copy

Copy

271414432

 

题意：

一棵树，n个节点，k种颜色，距离<=2的树不能冉同一种颜色，问方案数量？

分析：

距离为<=2的节点不能染相同的颜色,

我们首先假设一下这题只考虑一个距离为<=1的节点不能染相同的颜色,ans=k*（k-1）^(n-1)，即对于每一个点来说有几种可以选的颜色，如果父亲节点选了，选择有k种，则儿子节点不能选，只有（k-1）种，儿子的兄弟节点也只有相同的k-1种，儿子节点的儿子节点也有n-1种;

那么距离为<=2的叶子节点不能染相同的颜色,思想一样，考虑当前点可以选择的颜色个数。

如果父亲节点选了，选择有k种，则儿子节点不能选，只有（k-2）种，因为其父亲和祖父；

父亲的第二个孩子，只有相同的k-3种，因为其父亲和祖父和第一个孩子；

父亲的第三个孩子，只有相同的k-3种，因为其父亲和祖父和第一、二个孩子；

。。。。

但有一个坑点需要注意：当前节点为根，没有祖父，其孩子为k-1种选择

 

import sys

sys.setrecursionlimit(10**9)
#解决递归限制


N, K = map(int, input().split())
mod = 1000000007

def dfs(K, graph, now, source):
  if source == -1:
    can_use_color_num = K - 1
  else:
    can_use_color_num = K - 2
  if len(graph[now]) > K:
    return 0
  else:
    case_num = 1
    for dst in graph[now]:
      if dst == source:
        continue
      case_num *= can_use_color_num
      can_use_color_num -= 1
      case_num %= mod
    for dst in graph[now]:
      if dst == source:
        continue
      case_num *= dfs(K, graph, dst, now)
      case_num %= mod
    return case_num


graph = [[] for i in range(N)]
for i in range(N - 1):
  a, b = map(lambda x:int(x)-1, input().split())
  graph[a].append(b)
  graph[b].append(a)

answer = K * dfs(K, graph, 0, -1)
answer %= mod
print(answer)