Redaksiya

Let's start coloring the tree from vertex $1$. It can be colored with $k$ colors.

Suppose we are at vertex $v$. If it has no parent ($v=1$), then its children can be colored with $k-1$ colors. If $v$ has a parent, then its children can be colored with $k-2$ colors. Since the distance between the children of vertex $v$ is $2$, they cannot be colored with the same color.

If vertex $v$ has a parent, then the first child of $v$ can be colored with $k-2$ colors, the second child with $k-3$ 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 $6$ 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 $5$ can be colored with $2$ colors, as its color must not match the colors of vertices $1$ and $4$. The total number of ways to color the tree is equal to $4\cdot 3\cdot 2\cdot 1\cdot 2=48$.

Algorithm implementation

Store the input graph in the adjacency list $g$. Declare a constant $MOD$ 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 $v$ modulo $MOD=10^9+7$. The parent of vertex $v$ is $p$. Vertex $v$ can be colored with $col$ colors.

int dfs(int v, int col, int p = -1)
{
  long long res = col;

We are at vertex $v$. In the variable $cnt$, we keep track of the number of colors that cannot be used to paint the children of vertex $v$. Initially, set $cnt=1$, since the children of vertex $v$ cannot have the same color as vertex $v$. If vertex $v$ has a parent ($p\ne -1$), then the children of vertex $v$ also cannot have the color of $v$'s parent.

  int cnt = 1;
  if (p >= 0) cnt++;

Iterate over the sons $to$ of the vertex $v$.

  for (int to : g[v])
    if (to != p)
    {

The vertex $to$ can be colored with $k-cnt$ colors. With each son, the $cnt$ value will be increased by $1$.

      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 $1$. Vertex number $1$ can be colored with $k$ 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 $MOD$ for the modulo.

MOD = 1000000007

The dfs function returns the number of ways to color the tree rooted at vertex $v$ modulo $MOD=10^9+7$. The parent of vertex $v$ is $p$. Vertex $v$ can be colored with $col$ colors.

def dfs(v, col, p = -1):
  res = col

We are at vertex $v$. In the variable $cnt$, we keep track of the number of colors that cannot be used to paint the children of vertex $v$. Initially, set $cnt=1$, since the children of vertex $v$ cannot have the same color as vertex $v$. If vertex $v$ has a parent ($p\ne -1$), then the children of vertex $v$ also cannot have the color of $v$'s parent.

  cnt = 1
  if p >= 0: cnt += 1

Iterate over the sons $to$ of the vertex $v$.

  for to in g[v]:
    if to != p:

The vertex $to$ can be colored with $k-cnt$ colors. With each son, the $cnt$ value will be increased by $1$.

      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 $1$. Vertex number $1$ can be colored with $k$ colors.

res = dfs(1, k)

Print the answer.

print(res)