11609. Найти пары

Задано дерево с n вершинами. Вершина 1 является корнем. Между любыми двумя вершинами существует единственный путь. Определим d(i, j) как количество ребер на пути между вершинами i и j.

Найдите количество таких пар (i, j), что d(i, j) = d(i, 1) – d(j, 1).

Вход. Первая строка содержит количество вершин n (1 ≤ n ≤ 10⁵) в дереве. Каждая из следующих n – 1 строк задает ребро дерева.

Выход. Выведите количество таких пар (i, j), что d(i, j) = d(i, 1) – d(j, 1).

Пример входа	Пример выхода
5 1 2 2 3 2 4 4 5	13

РЕШЕНИЕ

графы – поиск в глубину

Анализ алгоритма

Условие 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, то вычисляем глубину to и запускаем из to поиск в глубину.

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