Відстань у дереві#4104

Відстань у дереві (Розбір)

Нехай $d p [v] [x]$ dp[v][x] — кількість вершин, що знаходяться на відстані $x$ від вершини $v$ в піддереві з коренем у $v$ . Покладемо $d p [v] [0] = 1$ dp[v][0] = 1, оскільки на відстані $0$ від вершини $v$ знаходиться тільки сама вершина $v$ .

Нехай вершина $v$ має синів $u_{1}, u_{2}, ..., u_{p}$ u_1, u_2, ..., u_p. Тоді кількість вершин, що знаходяться на відстані $x$ від вершини $v$ в піддереві з коренем у $v$ , дорівнює сумі кількості вершин, що знаходяться на відстані $x - 1$ x - 1 від кожної з вершин $u_{1}, u_{2}, ..., u_{p}$ u_1, u_2, ..., u_p.

d p [v] [x] = i = 1 \sum p d p [u_{i}] [x - 1]

dp[v][x] = \sum_{i = 1}^p dp[u_i][x-1]

Розглянемо піддерево з коренем у вершині $v$ . Порахуємо кількість шляхів довжини $k$ , повністю лежачих у цьому піддереві і проходящих через вершину $v$ . Розіб'ємо такі шляхи на дві групи:

шляхи, що починаються у вершині $v$ . Їх кількість дорівнює $d p [v] [k]$ dp[v][k].
шляхи, що проходять через вершину $v$ , кінці яких розташовані в піддеревах синів вершини $v$ .

Розглянемо шлях довжини $k$ , проходящий через корінь піддерева $v$ . Нехай відстань від вершини $v$ до одного з кінців шляху дорівнює $x (1 \leq x \leq k - 1)$ x~(1 \le x \le k - 1). Припустимо, що цей кінець знаходиться у піддереві певного сина $u_{1}$ u_1. Тоді відстань від вершини $v$ до другого кінця шляху дорівнює $k - x$ k - x, причому другий кінець може знаходитися в будь-якому іншому піддереві, відмінному від піддерева з коренем у $u_{1}$ u_1.

Якщо відстань від першого кінця шляху до вершини $v$ дорівнює $x$ , то відстань від цього кінця до вершини $u_{1}$ u_1 дорівнює $x - 1$ x - 1. Число способів вибрати такий перший кінець дорівнює $d p [u_{1}] [x - 1]$ dp[u_1][x - 1].

Число способів вибрати другий кінець шляху дорівнює кількості вершин, що знаходяться на відстані $k - x$ k - x від вершини $v$ , але не належать піддереву з коренем у $u_{1}$ u_1. Це кількість дорівнює

d p [v] [k - x] - d p [u_{1}] [k - x - 1]

dp[v][k - x] - dp[u_1][k - x - 1]

Згідно з правилом добутку, загальна кількість таких шляхів дорівнює

d p [u_{1}] [x - 1] \cdot (d p [v] [k - x] - d p [u_{1}] [k - x - 1])

dp[u_1][x - 1] \cdot (dp[v][k - x] - dp[u_1][k - x - 1])

Таким чином, загальна кількість таких шляхів дорівнює

d p [v] [k] + 1/2 i = 1 \sum p x = 1 \sum k - 1 d p [u_{i}] [x - 1] \cdot (d p [v] [k - x] - d p [u_{i}] [k - x - 1])

dp[v][k] + 1/2 \sum_{i = 1}^p \sum_{x = 1}^{k - 1} dp[u_i][x-1] \cdot (dp[v][k-x]-dp[u_i][k-x-1])

Залишається просумувати зазначене вираз по всім вершинам дерева.

Приклад

Розглянемо дерево, наведене в прикладі.

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

В масиві $g$ зберігається представлення дерева у вигляді списків суміжності. Масив $d p$ використовується для реалізації динамічного програмування.

vector<vector<int> > g;
vector<vector<int> > dp;

Функція dfs виконує пошук в глибину, починаючи з вершини $v$ . Значення $p$ позначає предка вершини $v$ . В процесі обходу будуємо масив $d p$ .

void dfs(int v, int p = -1)
{

Ініціалізуємо масив $d p [v]$ .

  dp[v].resize(k+1);
  dp[v][0] = 1;

Перебираємо всі вершини $t o$ , в які можна перейти з вершини $v$ .

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

    dfs(to,v);

    for(int x = 1; x <= k; x++)
      dp[v][x] += dp[to][x-1];
  }
}

C++

Функція dfsRes виконує пошук в глибину, починаючи з вершини $v$ , і призначена для обчислення шуканої відповіді.

void dfsRes(int v, int p = -1)
{

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

    dfsRes(to,v);

    for(int x = 1; x < k; x++)
      resA += 1LL * dp[to][x-1] * (dp[v][k-x] - dp[to][k-x-1]);
  }

  resB += dp[v][k];
}

C++

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

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

C++

Першим пошуком в глибину будуємо масив $d p$ .

dp.resize(n+1);
dfs(1);

Підраховуємо та виводимо шукану відповідь.

resA = resB = 0;
dfsRes(1);
res = resA / 2 + resB;
printf("%lld\n",res);

Python реалізація

Функція dfs виконує пошук в глибину, починаючи з вершини $v$ . Значення $p позначаєпредкавершини$ v $. Впроцесіобходубудуємомасив$ dp $. \begin{lstlisting}[language=Python] def dfs(v, p): \end{lstlisting} Ініціалізуємо масив$ dp[v] $. \begin{lstlisting}[language=Python] dp[v] = [0] * (k + 1) dp[v][0] = 1 \end{lstlisting} Перебираємо всі вершини$ to $, вякіможнаперейтизвершини$ v $. \begin{lstlisting}[language=Python] for to in g[v]: if to == p: continue dfs(to, v) for x in range(1, k + 1): dp[v][x] += dp[to][x - 1] \end{lstlisting} Функція \textbf{dfsRes} виконує пошук в глибину, починаючи з вершини$ . \begin{lstlisting}[language=Python] for to in g[v]: if to == p: continue dfs(to, v) for x in range(1, k + 1): dp[v][x] += dp[to][x - 1] \end{lstlisting} Функція \textbf{dfsRes} виконує пошук в глибину, починаючи з вершини v $, і призначена для обчислення шуканої відповіді. \begin{lstlisting}[language=Python] def dfsRes(v, p): global resA, resB for to in g[v]: if to == p: continue dfsRes(to, v) for x in range(1, k): resA += dp[to][x - 1] * (dp[v][k - x] - dp[to][k - x - 1]) resB += dp[v][k] \end{lstlisting} Основна частина програми. Читаємо кількість вершин дерева$ , і призначена для обчислення шуканої відповіді. \begin{lstlisting}[language=Python] def dfsRes(v, p): global resA, resB for to in g[v]: if to == p: continue dfsRes(to, v) for x in range(1, k): resA += dp[to][x - 1] * (dp[v][k - x] - dp[to][k - x - 1]) resB += dp[v][k] \end{lstlisting} Основна частина програми. Читаємо кількість вершин дерева n $тавідстаньміжвершинами$ k $. \begin{lstlisting}[language=Python] n, k = map(int, input().split()) \end{lstlisting} У списку$ g $зберігаєтьсяпредставленнядереваувиглядісписківсуміжності . Список$ dp $використовується для реалізації динамічного програмування. \begin{lstlisting}[language=Python] g = [[] for _ in range(n + 1)] dp = [[] for _ in range(n + 1)] \end{lstlisting} Читаємо список ребер. Будуємо дерево. \begin{lstlisting}[language=Python] for _ in range(n - 1): a, b = map(int, input().split()) g[a].append(b) g[b].append(a) \end{lstlisting} Першим пошуком в глибину будуємо масив$ використовується для реалізації динамічного програмування. \begin{lstlisting}[language=Python] g = [[] for _ in range(n + 1)] dp = [[] for _ in range(n + 1)] \end{lstlisting} Читаємо список ребер. Будуємо дерево. \begin{lstlisting}[language=Python] for _ in range(n - 1): a, b = map(int, input().split()) g[a].append(b) g[b].append(a) \end{lstlisting} Першим пошуком в глибину будуємо масив dp$.

dfs(1, -1)

Підраховуємо та виводимо шукану відповідь.

resA, resB = 0, 0
dfsRes(1, -1)
res = resA // 2 + resB
print(res)