Вірусне дерево 2 (Розбір)

Почнемо розфарбовку дерева з вершини $1$. Її можна пофарбувати $k$ кольорами.

Нехай ми знаходимося у вершині $v$. Якщо вона не має предка (v = 1), то синів можна розфарбувати k - 1 кольорами. Якщо вершина $v$ має предка, то її синів можна пофарбувати k - 2 кольорами. Оскільки відстань між синами вершини $v$ дорівнює $2$, то їх не можна фарбувати однаковими кольорами.

Нехай для визначеності вершина $v$ має предка. Тоді першого сина $v$ можна пофарбувати k - 2 кольорами, другого сина k - 3 кольорами і так далі.

Залишається перемножити кількості кольорів, якими можна розфарбувати кожну вершину.

Приклад

Для першого прикладу маємо $6$ способів розфарбування:

Розглянемо другий тест. Біля кожної вершини вкажемо кількість кольорів, якими її можна пофарбувати. Наприклад, вершину $5$ можна пофарбувати 2 кольорами, оскільки її колір не повинен збігатися з кольорами вершин $1$ і $4$. Загальна кількість способів пофарбувати дерево дорівнює 4 \cdot 3 \cdot 2 \cdot 1 \cdot 2 = 48.

Реалізація алгоритму

Вхідний граф зберігаємо в списку суміжності $g$. Оголосимо константу — модуль $MOD$.

#define MOD 1000000007
vector<vector<int>> g;

Функція dfs повертає кількість способів розфарбувати дерево з коренем у вершині $v$ за модулем MOD = 10^9 + 7. Предком вершини $v$ є $p$. Вершину $v$ можна пофарбувати $col$ кольорами.

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

Ми знаходимося у вершині $v$. У змінній $cnt$ відзначаємо кількість кольорів, якими не можна фарбувати синів вершини $v$. Спочатку встановимо cnt = 1, оскільки сини вершини $v$ не можуть мати колір вершини $v$. Якщо вершина $v$ має предка (p \ne -1), то сини вершини $v$ також не можуть мати колір предка $v$.

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

Перебираємо синів to вершини $v$.

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

Вершина $to$ може бути пофарбована k - cnt кольорами. З кожним сином значення $cnt$ буде збільшуватись на $1$.

      res = (res * dfs(to, k - cnt, v)) % MOD;
      cnt++;
    }
  return res;
}

Основна частина програми. Читаємо вхідні дані.

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);
}

Запускаємо пошук у глибину з вершини $1$. Вершину номер $1$ можна пофарбувати $k$ кольорами.

res = dfs(1, k);

Виводимо відповідь.

printf("%lld\n", res);

Python реалізація

Встановимо максимальну глибину стека інтерпретатора Python.

import sys
sys.setrecursionlimit(1000000)

Оголосимо константу — модуль $MOD$.

MOD = 1000000007

Функція dfs повертає кількість способів розфарбувати дерево з коренем у вершині $v$ за модулем MOD = 10^9 + 7. Предком вершини $v$ є $p$. Вершину $v$ можна пофарбувати $col$ кольорами.

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

Ми знаходимося у вершині $v$. У змінній $cnt$ відзначаємо кількість кольорів, якими не можна фарбувати синів вершини $v$. Спочатку встановимо cnt = 1, оскільки сини вершини $v$ не можуть мати колір вершини $v$. Якщо вершина $v$ має предка (p \ne -1), тоді сини вершини $v$ також не можуть мати колір предка $v$.

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

Перебираємо синів $to$ вершини $v$.

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

Вершина $to$ може бути пофарбована k - cnt кольорами. З кожним сином значення $cnt$ буде збільшуватись на $1$.

      res = (res * dfs(to, k - cnt, v)) % MOD
      cnt += 1
  return res

Основна частина програми. Читаємо вхідні дані.

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)

Запускаємо пошук у глибину з вершини $1$. Вершину номер $1$ можна пофарбувати $k$ кольорами.

res = dfs(1, k)

Виводимо відповідь.

print(res)