Editorial
Let's start coloring the tree from vertex . It can be colored with colors.
Suppose we are at vertex . If it has no parent (), then its children can be colored with colors. If has a parent, then its children can be colored with colors. Since the distance between the children of vertex is , they cannot be colored with the same color.
If vertex has a parent, then the first child of can be colored with colors, the second child with colors, and so on.
It remains to multiply the numbers of colors that can be used to color each vertex.
Example
For the first example, there are ways of coloring:
Let's consider the second test. Next to each vertex, we'll indicate the number of colors it can be colored with. For example, vertex can be colored with colors, as its color must not match the colors of vertices and . The total number of ways to color the tree is equal to .
Algorithm implementation
Store the input graph in the adjacency list . Declare a constant for the modulo.
#define MOD 1000000007 vector<vector<int>> g;
The dfs function returns the number of ways to color the tree rooted at vertex modulo . The parent of vertex is . Vertex can be colored with colors.
int dfs(int v, int col, int p = -1) { long long res = col;
We are at vertex . In the variable , we keep track of the number of colors that cannot be used to paint the children of vertex . Initially, set , since the children of vertex cannot have the same color as vertex . If vertex has a parent (), then the children of vertex also cannot have the color of 's parent.
int cnt = 1; if (p >= 0) cnt++;
Iterate over the sons of the vertex .
for (int to : g[v]) if (to != p) {
The vertex can be colored with colors. With each son, the value will be increased by .
res = (res * dfs(to, k - cnt, v)) % MOD; cnt++; } return res; }
The main part of the program. Read the input data.
scanf("%d %d", &n, &k); g.resize(n + 1); for (i = 0; i < n - 1; i++) { scanf("%d %d", &a, &b); g[a].push_back(b); g[b].push_back(a); }
Start the depth-first search from vertex . Vertex number can be colored with colors.
res = dfs(1, k);
Print the answer.
printf("%lld\n", res);
Python implementation
Set the maximum depth of the Python interpreter stack to the specified limit.
import sys sys.setrecursionlimit(1000000)
Declare a constant for the modulo.
MOD = 1000000007
The dfs function returns the number of ways to color the tree rooted at vertex modulo . The parent of vertex is . Vertex can be colored with colors.
def dfs(v, col, p = -1): res = col
We are at vertex . In the variable , we keep track of the number of colors that cannot be used to paint the children of vertex . Initially, set , since the children of vertex cannot have the same color as vertex . If vertex has a parent (), then the children of vertex also cannot have the color of 's parent.
cnt = 1 if p >= 0: cnt += 1
Iterate over the sons of the vertex .
for to in g[v]: if to != p:
The vertex can be colored with colors. With each son, the value will be increased by .
res = (res * dfs(to, k - cnt, v)) % MOD cnt += 1 return res
The main part of the program. Read the input data.
n, k = map(int, input().split()) g = [[] for _ in range(n + 1)] for _ in range(n - 1): a, b = map(int, input().split()) g[a].append(b) g[b].append(a)
Start the depth-first search from vertex . Vertex number can be colored with colors.
res = dfs(1, k)
Print the answer.
print(res)