4104. Расстояние в дереве
Деревом
называется связный граф, не содержащий циклов.
Расстоянием
между двумя вершинами дерева называется длина (в ребрах) кратчайшего пути между
этими вершинами.
Дано дерево из n вершин и положительное число k. Посчитайте количество различных пар
вершин дерева, расстояние между которыми равно k. Обратите внимание, что пары (v,
u) и (u, v) считаются одной и
той же парой.
Вход. В первой строке записаны два целых числа n и k
(1 ≤ n ≤ 50000, 1 ≤
k ≤ 500) – количество вершин
дерева и требуемое расстояние между вершинами.
В следующих n – 1 строках записаны ребра дерева в
формате ai bi (1 ≤ ai, bi ≤ n, ai ≠ bi), где ai и bi
– вершины дерева, соединенные i-ым
ребром. Все заданные ребра различны.
Выход. Выведите единственное целое число – количество различных
пар вершин дерева, расстояние между которыми равно k.
|
Пример
входа |
Пример
выхода |
|
5 2 1 2 2 3 3 4 2 5 |
4 |
РЕШЕНИЕ
динамическое
программирование на дереве
Анализ алгоритма
Пусть dp[v][x] – количество
вершин на расстоянии x от вершины v в поддереве с корнем v. Положим dp[v][0] = 1, считая что на расстоянии 0 от вершины v находится только одна вершина v. Пусть вершина v имеет сыновей u1,
u2, .., up. Тогда
dp[v][x] = 
Количество
вершин, находящихся на расстоянии x от вершины v (поддереве с корнем v), равно количеству вершин,
находящихся от сыновей вершины v на
расстоянии x – 1.
Рассмотрим
поддерево с корнем v. Подсчитаем,
сколько существует в нем путей длины k, которые
проходят через вершину v. Разобъем их на две группы:
·
пути, которые начинаются в v. Их количество
равно dp[v][k].
·
пути, которые проходят через v концы которых
лежат в поддеревьях сыновей v.
Рассмотрим путь
длины k, проходящего через корень
поддерева v. Пусть расстояние от v до одного из его концов равно x, 1 ≤ x ≤ k – 1 (и например, этот
конец расположен в поддереве сына u1) тогда
расстояние от v до другого его конца равно k – x (второй конец
может располагаться в любом другом поддереве, отличном от поддерева с корнем в u1).
Если расстояние
от первого конца пути до v равно x, то расстояние от него до u1 равно x – 1. Количество способов, которыми можно выбрать этот конец,
равно dp[u1][x – 1]. Количество способов, которыми
можно выбрать второй конец пути, равно количеству вершин, расположенных на расстоянии k – x от v, но не лежащих в поддереве с корнем u1. Их число равно dp[v][k – x] – dp[u1][k – x – 1]. Согласно правилу произведения
количество таких путей равно dp[u1][x – 1] * (dp[v][k – x] – dp[u1][k – x – 1]).
Таким образом их
количество равно
dp[v][k] + 
Остается
просуммировать указанное выражение по всем вершинам.
Пример
Рассмотрим
дерево, приведенное в примере.
Реализация алгоритма
В массиве g содержим дерево.
Массив dp используется
для динамики.
vector<vector<int> > g;
vector<vector<int> > dp;
Функция dfs реализует поиск в глубину из вершины v. Значение p является
предком v. Строим массив dp.
void dfs(int
v, int p = -1)
{
Инициализируем массив dp[v].
dp[v].resize(k+1);
dp[v][0] = 1;
for(int i = 0; i < g[v].size(); i++)
{
Перебираем все вершины to, куда можно
попасть из v.
int to =
g[v][i];
if (to ==
p) continue;
dfs(to,v);
for(int x = 1; x <= k; x++)
dp[v][x] += dp[to][x-1];
}
}
Функция dfsRes реализует поиск в глубину из вершины v с вычислением ответа.
void dfsRes(int
v, int p = -1)
{
for(int i = 0; i < g[v].size(); i++)
{
int to =
g[v][i];
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];
}
Основная часть программы. Читаем входные данные. Строим
дерево.
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);
}
Первым поиском в глубину строим массив dp.
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.
def dfs(v, p):
Инициализируем массив dp[v].
dp[v] = [0] * (k + 1)
dp[v][0] = 1
Перебираем все вершины to, куда можно
попасть из v.
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]
Функция dfsRes реализует поиск в глубину из вершины v с вычислением ответа.
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]
Основная часть программы. Читаем количество вершин дерева n и расстояние
между вершинами k.
n, k = map(int, input().split())
В списке g
содержим дерево. Список dp используется
для динамики.
g = [[] for _ in
range(n + 1)]
dp = [[] for _ in
range(n + 1)]
Читаем список ребер. Строим дерево.
for _ in range(n - 1):
a, b = map(int, input().split())
g[a].append(b)
g[b].append(a)
Первым поиском в глубину строим массив dp.
dfs(1, -1)
Подсчитываем и выводим ответ.
resA, resB = 0, 0
dfsRes(1, -1)
res = resA // 2 + resB
print(res)