Знайти пари (Розбір)

Умова d(i, j) = d(i, 1) - d(j, 1) виконується для кожної пари вершин $(i,j)$, у якої $j$ є предком $i$ при пошуку в глибину з вершини $1$. Розглянемо приклад:

• d(4, 1) = 2, d(4, 1) - d(1, 1) = 2 - 0 = 2;

• d(6, 3) = 2, d(6, 1) - d(3, 1) = 3 - 1 = 2;

• d(2, 2) = 0, d(2, 1) - d(2, 1) = 1 - 1 = 0;

• d(6, 1) = 3, d(6, 1) - d(1, 1) = 3 - 0 = 3;

Глибиною $h[v]$ вершини $v$ назвемо кількість вершин на шляху від кореня (вершини $1$) до $v$. На малюнку глибина вказана біля кожної вершини.

Зафіксуємо вершину $i$ і відповімо на питання: скільки існує вершин $j$, для яких виконується умова d(i, j) = d(i, 1) - d(j, 1).

Наприклад, для i = 6 існують $4$ такі вершини: j = 1, 3, 4, 6. Зазначимо, що h[6] = 4. Таким чином, для фіксованої вершини $i$ існує рівно $h[i]$ підходящих вершин $j$.

Для розв'язання задачі необхідно для кожної вершини $v$ дерева обчислити її глибину $h[v]$, а потім знайти суму всіх значень глибин по дереву.

Приклад

Розглянемо граф з прикладу. Біля кожної вершини запишемо її глибину.

Сума глибин всіх вершин дорівнює 1 + 2 + 3 + 3 + 4 = 13.

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

Оголосимо список суміжності графа $g$. Глибину вершин зберігаємо в масиві $h$.

vector<vector<int> > g;
vector<int> h;

Функція dfs для кожної вершини $v$ обчислює її глибину. Предком вершини $v$ є вершина $p$.

int dfs(int v, int p)
{
  for (int to : g[v])

Розглядаємо ребро (v, to). Якщо вершина $to$ не є предком вершини $v$, обчислюємо її глибину і запускаємо з неї пошук в глибину.

    if (to != p)
    {
      h[to] = h[v] + 1;
      dfs(to, v);
    }
  return h[v];
}

Основна частина програми. Зчитуємо кількість вершин дерева $n$.

scanf("%d", &n);
g.resize(n + 1);

Будуємо дерево.

for (i = 2; i <= n; i++)
{
  scanf("%d %d", &a, &b);
  g[a].push_back(b);
  g[b].push_back(a);
}

Глибину вершини $1$ покладемо рівною $1$.

h.resize(n + 1);
h[1] = 1;

Запускаємо пошук в глибину, починаючи з вершини $1$.

dfs(1, -1);

Обчислюємо відповідь — суму глибин всіх вершин.

res = 0;
for (i = 1; i <= n; i++)
  res += h[i];

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

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