DP на дереве с накоплением подответов — задача «Distance in Tree» РАЗБОР

DP на дереве — естественное обобщение линейного DP: вместо того чтобы идти по массиву слева направо, обход идёт от листьев к корню, и в каждой вершине ответ для поддерева складывается из ответов для поддеревьев детей. Эта статья разбирает классический шаблон такого DP с накоплением подответов через свёртку — на задаче, где этот приём раскрывается особенно прозрачно.

Идея, к которой стоит привыкнуть: при обходе дерева в каждой вершине мы держим не одно число, а массив длины $k+1$, где $d$-я ячейка отвечает на «сколько вершин в поддереве на расстоянии $d$ от корня поддерева». Когда вершина сшивает в себя очередную дочернюю вершину, мы сначала подсчитываем пары длины $k$, у которых эта вершина — точка поворота, и только потом вливаем массив дочерней в свой. Этот «свёртка перед слиянием» — главный приём шаблона.

---

Что дано

Дано неориентированное дерево с $n$ вершинами (нумерация $1, 2, \ldots, n$) и целое число $k$. Нужно найти количество неупорядоченных пар вершин $\{u, v\}$, $u \ne v$, таких, что длина простого пути от $u$ до $v$ в дереве (число рёбер на пути) равна ровно $k$.

Ограничения

• $1 \le n \le 50\,000$.
• $1 \le k \le 500$.
• Время — 2 секунды, память — 256 МБ.

Формат ввода

n k
u_1 v_1
u_2 v_2
...
u_{n-1} v_{n-1}

Первая строка — два числа: количество вершин и требуемое расстояние. Затем $n-1$ строка, каждая с парой вершин — концами очередного ребра.

Формат вывода

Одно целое число — количество искомых пар.

Пример 1

Вход:

5 2
1 2
2 3
3 4
2 5

Вывод:

4

Дерево выглядит так:

1
└── 2
    ├── 3
    │   └── 4
    └── 5

Пары на расстоянии 2: $\{1, 3\}$ (через вершину 2), $\{1, 5\}$ (через 2), $\{3, 5\}$ (через 2), $\{2, 4\}$ (через 3). Итого 4.

Пример 2

Вход:

5 3
1 2
2 3
3 4
4 5

Вывод:

2

Дерево — вытянутая в линию цепочка $1 - 2 - 3 - 4 - 5$. Пары на расстоянии 3: $\{1, 4\}$ и $\{2, 5\}$. Итого 2.

Разберём первый пример руками

Подвесим дерево за вершину 1 как корень. У 1 одна дочерняя вершина — 2; у 2 две дочерние вершины — 3 и 5; у 3 одна дочерняя вершина — 4; у 4 и 5 дочерних вершин нет.

Подумаем, как все 4 искомые пары устроены относительно этого корня:

Каждая пара $\{u, v\}$ имеет ровно одну общую «точку поворота» — наименьшего общего предка (LCA). Если мы научимся для каждой вершины $w$ считать «сколько пар $\{u, v\}$ имеют расстояние $k$ и LCA в $w$» — можно просто просуммировать по всем $w$, и общая сумма даст ответ без двойного счёта.

Считаем для $w$ как для LCA: расстояние $u \to v = (\text{расстояние от } u \text{ до } w) + (\text{расстояние от } v \text{ до } w)$. Если $u$ — это сам $w$, то первое слагаемое — 0.

Если в каждой вершине $w$ запомнить таблицу «сколько вершин в поддереве $w$ находятся на расстоянии $d$ от $w$» (для $d$ от 0 до $k$) — посчитать пары с LCA в $w$ можно за $O(k)$: перебрать пару $(d_1, d_2)$ с $d_1 + d_2 = k$ и $u, v$ из разных поддеревьев $w$ (или одно из них — сам $w$). Алгоритм нашего шаблона делает ровно это.

---

Идея решения: tree DP со свёрткой через точку поворота

Алгоритм в одну фразу.

Подвешиваем дерево за корень, делаем DFS. В каждой вершине $v$ держим массив $\mathrm{dp}[v]$ длины $k+1$, где $\mathrm{dp}[v][d]$ — количество вершин в поддереве с корнем $v$ на расстоянии $d$ от $v$. Сшивая поддеревья дочерних вершин, перед слиянием очередной дочерней вершины $c$ в $\mathrm{dp}[v]$ добавляем к ответу свёртку $\sum_{i+j=k-1} \mathrm{dp}[v][i] \cdot \mathrm{dp}[c][j]$ — это пары длины $k$, у которых $v$ — точка поворота.

Почему это работает

Зафиксируем произвольный корень — пусть это вершина 1. Любая пара $\{u, v\}$ имеет в этом корневом дереве уникальный наименьший общий предок $w$ — то есть путь от $u$ до $v$ идёт вверх до $w$, а потом вниз. Это даёт ключевую декомпозицию: общий ответ равен сумме по всем вершинам $w$ количества пар длины $k$, у которых $w$ — это LCA. Считать слагаемые будем по ходу DFS — когда обрабатываем $w$ и сшиваем её поддеревья.

Сшиваем поддерево очередной дочерней вершины $c$ в уже накопленный $\mathrm{dp}[w]$ (который пока описывает саму $w$ и все её ранее обработанные поддеревья). Любая пара $\{u, v\}$ с LCA $= w$, где $u$ из «уже накопленного», а $v$ — из поддерева $c$, выражается через расстояния так: расстояние $u \to w$ — это какое-то $i$, расстояние $w \to v$ — это $1 + j$, где $j$ — расстояние от $c$ до $v$ (плюс единица за ребро $w \leftrightarrow c$). Итог $i + 1 + j = k$, то есть $j = k - 1 - i$.

Считая такие пары по всем возможным $i$ — приходим к свёртке:

$$\Delta_w(c) = \sum_{i = 0}^{k - 1} \mathrm{dp}[w][i] \cdot \mathrm{dp}[c][k - 1 - i].$$

Прибавляем $\Delta_w(c)$ к ответу — и сливаем $c$ в $w$ со сдвигом $+1$: dp[w][d+1] += dp[c][d] для всех $d$ (вершины из поддерева $c$ удалены от $w$ на единицу больше, чем от $c$). Когда переходим к следующей дочерней вершине $w$, в $\mathrm{dp}[w]$ уже учтены и сама $w$, и все ранее обработанные поддеревья — следующая свёртка автоматически считает пары «между новым поддеревом и всем уже сшитым».

Почему пары не дублируются и не теряются

• Не дублируются: пара $\{u, v\}$ имеет уникальный LCA $w$, и алгоритм добавляет её в ответ ровно один раз — в момент, когда обрабатывается первый из двух поддеревьев $w$, содержащий один из концов пары (то есть в момент сшивания второго поддерева, содержащего второй конец).
• Не теряются: все пары попадают в одну из категорий. Если $u = w$ (одна из вершин — сам LCA), то $u$ присутствует в $\mathrm{dp}[w][0]$ ещё до начала перебора детей; пара учтётся при сшивании поддерева, содержащего $v$.

Эти два свойства дают корректность алгоритма без дополнительных проверок. Ошибочные варианты «перебрать пары вершин и посчитать расстояние через LCA» имеют сложность $O(n^2)$ и на $n = 50\,000$ не проходят по времени.

---

Решение: псевдокод

# глобально
ans ← 0
dp[v] для каждого v — массив длины k+1, изначально пустой

функция dfs(v, parent):
    dp[v] ← массив [0, 0, ..., 0] длины k+1
    dp[v][0] ← 1                         # сама вершина v на расстоянии 0
    для каждого соседа c вершины v, c ≠ parent:
        dfs(c, v)
        # ① подсчёт пар с LCA в v: один конец в уже сшитом dp[v], второй — в dp[c]
        для i от 0 до k-1:
            ans ← ans + dp[v][i] * dp[c][k - 1 - i]
        # ② слияние dp[c] в dp[v] со сдвигом +1
        для d от 0 до k-1:
            dp[v][d + 1] ← dp[v][d + 1] + dp[c][d]
        освободить dp[c]                  # больше не понадобится, экономим память

# запуск
dfs(1, 0)
вывести ans

Почему важен именно такой порядок «① свёртка → ② слияние». Если поменять местами, после слияния в $\mathrm{dp}[v]$ окажутся вершины из $T_c$, и при «свёртке после» мы будем считать в том числе пары вида $(u \in T_c, v \in T_c)$ — но они не имеют LCA в $v$, их LCA лежит глубже, в самом $c$ или ниже. Эти пары уже учтены раньше (в момент, когда $c$ сшивал свои собственные поддеревья). Получится двойной счёт.

Сложность: каждое ребро дерева вносит $O(k)$ работы (два цикла длины $k$). Рёбер $n - 1$. Итог — $O(n \cdot k)$ по времени.

---

Код решения

Переключи вкладку, если хочешь видеть второй язык. Для Python важная деталь: рекурсия глубиной до $50\,000$ упирается в дефолтный лимит (~1000) и стек ОС, поэтому используем итеративный DFS с явным построением post-order. В C++ обычная рекурсия проходит — стек ОС обычно $8$ МБ.

Комментарии по реализации.

• Long long обязателен. Произведение $\mathrm{dp}[v][i] \cdot \mathrm{dp}[c][k - 1 - i]$ достигает $n^2 / 4 \approx 6 \cdot 10^8$, а сумма ответа — до $\binom{n}{2} \approx 1.25 \cdot 10^9$. В C++ — long long, иначе переполнение. В Python целые автоматически большие, заботы нет.
• Освобождение памяти. Без dp[c] = None (Python) или vector<long long>().swap(dp[c]) (C++) в худшем случае все $n$ массивов длины $k+1$ держатся одновременно — $\sim 200$ МБ. С освобождением одновременно живут только массивы вершин на «активной ветке» от корня до текущей.
• Рекурсия в Python — рискованно. Дефолтный лимит $1000$ ловит цепочку почти сразу. Поднятие через sys.setrecursionlimit(60000) помогает, но стек ОС при глубине $50\,000$ всё равно может переполниться. Итеративный DFS — надёжнее.
• Производительность Python. Для $n = 50\,000$, $k = 500$ — $2.5 \cdot 10^7$ операций в чистом Python. На жёстких лимитах ($1$–$2$ секунды) может не пройти. Для CF 161D проще писать на C++ или PyPy; для Python — векторизация через numpy или array.array('q') помогает.

---

Проверим на всех примерах из условия

Прогон выполняем по алгоритму выше, с подвешиванием за вершину 1. Будем записывать $\mathrm{dp}[v]$ как список (нумерация позиций — расстояние от $v$ до её потомка).

Пример 1: $n = 5, k = 2$, рёбра $1{-}2, 2{-}3, 3{-}4, 2{-}5$

Корень = 1. Post-order обхода: $4, 3, 5, 2, 1$.

dfs(4): $\mathrm{dp}[4] = [1, 0, 0]$. Детей нет, выход.

dfs(3): старт $\mathrm{dp}[3] = [1, 0, 0]$. Дочерняя вершина — 4.

• Свёртка: $i = 0$: $1 \cdot \mathrm{dp}[4][1] = 1 \cdot 0 = 0$. $i = 1$: $0 \cdot \mathrm{dp}[4][0] = 0$. $\Delta = 0$.
• Слияние: $\mathrm{dp}[3][1] \mathrel{+}= 1 \Rightarrow 1$. $\mathrm{dp}[3][2] \mathrel{+}= 0 \Rightarrow 0$.
• Итого $\mathrm{dp}[3] = [1, 1, 0]$. ans = 0.

dfs(5): $\mathrm{dp}[5] = [1, 0, 0]$. Детей нет.

dfs(2): старт $\mathrm{dp}[2] = [1, 0, 0]$. Дочерняя вершина 3 (первый из обработанных порядком).

• Свёртка: $i = 0$: $1 \cdot \mathrm{dp}[3][1] = 1 \cdot 1 = 1$. $i = 1$: $0 \cdot \mathrm{dp}[3][0] = 0$. $\Delta = 1$. ans = 1.
• Слияние: $\mathrm{dp}[2][1] \mathrel{+}= 1 \Rightarrow 1$. $\mathrm{dp}[2][2] \mathrel{+}= 1 \Rightarrow 1$.
• $\mathrm{dp}[2] = [1, 1, 1]$.

Дочерняя вершина 5.

• Свёртка: $i = 0$: $1 \cdot \mathrm{dp}[5][1] = 0$. $i = 1$: $1 \cdot \mathrm{dp}[5][0] = 1$. $\Delta = 1$. ans = 2.
• Слияние: $\mathrm{dp}[2][1] \mathrel{+}= 1 \Rightarrow 2$. $\mathrm{dp}[2][2] \mathrel{+}= 0 \Rightarrow 1$.
• Итого $\mathrm{dp}[2] = [1, 2, 1]$.

dfs(1): старт $\mathrm{dp}[1] = [1, 0, 0]$. Дочерняя вершина 2.

• Свёртка: $i = 0$: $1 \cdot \mathrm{dp}[2][1] = 1 \cdot 2 = 2$. $i = 1$: $0 \cdot \mathrm{dp}[2][0] = 0$. $\Delta = 2$. ans = 4.
• Слияние: $\mathrm{dp}[1][1] \mathrel{+}= 1 \Rightarrow 1$. $\mathrm{dp}[1][2] \mathrel{+}= 2 \Rightarrow 2$.

Ответ: 4 ✓.

Пример 2: $n = 5, k = 3$, рёбра $1{-}2, 2{-}3, 3{-}4, 4{-}5$ (линия)

Корень = 1. Post-order: $5, 4, 3, 2, 1$.

dfs(5): $\mathrm{dp}[5] = [1, 0, 0, 0]$. Нет детей.

dfs(4): старт $\mathrm{dp}[4] = [1, 0, 0, 0]$. Дочерняя вершина 5.

• Свёртка: $i = 0..2$: все слагаемые $\cdot \mathrm{dp}[5][2..0] = 1 \cdot 0 + 0 \cdot 0 + 0 \cdot 1 = 0$. $\Delta = 0$.
• Слияние: $\mathrm{dp}[4][1] \mathrel{+}= 1 \Rightarrow 1$. Остальные $\mathrm{dp}[5][1..2] = 0$ — без изменений.
• $\mathrm{dp}[4] = [1, 1, 0, 0]$.

dfs(3): старт $\mathrm{dp}[3] = [1, 0, 0, 0]$. Дочерняя вершина 4.

• Свёртка: $i = 0$: $1 \cdot \mathrm{dp}[4][2] = 0$. $i = 1$: $0 \cdot 1 = 0$. $i = 2$: $0 \cdot 1 = 0$. $\Delta = 0$.
• Слияние: $\mathrm{dp}[3][1] \mathrel{+}= 1 \Rightarrow 1$. $\mathrm{dp}[3][2] \mathrel{+}= 1 \Rightarrow 1$.
• $\mathrm{dp}[3] = [1, 1, 1, 0]$.

dfs(2): старт $\mathrm{dp}[2] = [1, 0, 0, 0]$. Дочерняя вершина 3.

• Свёртка: $i = 0$: $1 \cdot \mathrm{dp}[3][2] = 1 \cdot 1 = 1$. $i = 1$: $0$. $i = 2$: $0$. $\Delta = 1$. ans = 1.
• Слияние: $\mathrm{dp}[2][1] \mathrel{+}= 1 \Rightarrow 1$. $\mathrm{dp}[2][2] \mathrel{+}= 1 \Rightarrow 1$. $\mathrm{dp}[2][3] \mathrel{+}= 1 \Rightarrow 1$.
• $\mathrm{dp}[2] = [1, 1, 1, 1]$.

dfs(1): старт $\mathrm{dp}[1] = [1, 0, 0, 0]$. Дочерняя вершина 2.

• Свёртка: $i = 0$: $1 \cdot \mathrm{dp}[2][2] = 1 \cdot 1 = 1$. $i = 1$: $0$. $i = 2$: $0$. $\Delta = 1$. ans = 2.

Ответ: 2 ✓.

Оба примера сходятся с эталоном.

---

Крайние случаи

1. $n = 1$

Дерево — одна вершина, рёбер нет. $\mathrm{dp}[1] = [1, 0, \ldots, 0]$, циклов по детям не выполняется. Ответ: 0 для любого $k$. По условию $k \ge 1$ — пар $\{u, v\}$ с $u \ne v$ нет, всё корректно.

2. Звезда — одна центральная вершина и $n - 1$ листьев

$\mathrm{dp}[\text{лист}] = [1, 0, \ldots]$. В корне последовательно сшиваются $n - 1$ дочерняя вершина.

При $k = 2$: на каждой дочерней вершине свёртка даёт $\mathrm{dp}[\text{корень}][0] \cdot \mathrm{dp}[\text{лист}][1] + \mathrm{dp}[\text{корень}][1] \cdot \mathrm{dp}[\text{лист}][0]$. Первое слагаемое всегда 0 (у листа нет вершин на расстоянии 1). Второе — равно $\mathrm{dp}[\text{корень}][1]$, то есть числу уже обработанных листьев. Итог по всем дочерним вершинам — $0 + 1 + 2 + \ldots + (n - 2) = \binom{n - 1}{2}$. Для $n = 50\,000$: $\binom{49\,999}{2} \approx 1.25 \cdot 10^9$.

Этот случай — главная причина, почему ответ нужно держать в long long: $1.25 \cdot 10^9 > 2^{31} - 1 \approx 2.15 \cdot 10^9$ (всё ещё в int32, но при $k > 2$ ответ может ещё вырасти).

3. Линейное дерево («гусеница без листьев») — цепочка $1 - 2 - \ldots - n$

Глубина рекурсии в наивной C++-реализации — $n = 50\,000$. Стек ОС в $8$ МБ обычно проходит, но рекурсия с большими локальными переменными (вектор $\mathrm{dp}[v]$ длины $k + 1$) рискованна. На жёсткой границе по памяти переключиться на итеративный DFS — как в Python-коде выше.

В Python наивная рекурсия здесь падает с RecursionError на дефолтном лимите. Поднятие лимита до $60\,000$ + увеличение размера стека (threading.stack_size) помогает, но снова — итеративная версия надёжнее.

4. $k = 1$ — соседи

Свёртка с $k = 1$: $i$ ходит от 0 до 0, и $\Delta = \mathrm{dp}[v][0] \cdot \mathrm{dp}[c][0] = 1 \cdot 1 = 1$. Каждое ребро добавляет ровно 1 к ответу. Ответ — это число рёбер, то есть $n - 1$. Алгоритм сходится тривиально.

5. Большое $k$ — рассмотрим $k = n - 1$

Если в дереве есть пара вершин с расстоянием $n - 1$ — это значит, дерево содержит гамильтонов путь, а это бывает только для линейного дерева (цепочки). В этом случае ответ — 1 (две концевые вершины) или 0 (если дерево не цепочка). Алгоритм считает корректно. По ограничениям $k \le 500$, поэтому случай $k$ близкого к $n$ не встречается; но проверить алгоритм на нём всё равно полезно — выход за пределы массива нигде не происходит, так как индексы $[0, k]$ всегда валидны.

6. Память при $n \cdot k = 2.5 \cdot 10^7$

Без освобождения $\mathrm{dp}[c]$ после слияния все $n$ массивов длины $k + 1$ из long long ($8$ байт) держатся одновременно — $200$ МБ. Лимит задачи — $256$ МБ, формально проходим, но в Python из-за overhead'а каждого int (28 байт) запас исчезает. С освобождением живых одновременно — только массивы вершин на пути от корня до текущей: от $O(\log n)$ для сбалансированного дерева до $O(n)$ для цепочки.

---

Типичные ошибки

1. Перепутан порядок «свёртка → слияние». Если сначала влить $\mathrm{dp}[c]$ в $\mathrm{dp}[v]$, а потом считать свёртку, получится двойной счёт пар, которые лежат полностью в $T_c$ — их LCA не $v$, они уже учтены глубже.
2. Свёртка $i + j = k$ вместо $i + j = k - 1$. Самая частая алгебраическая ошибка: забыть про «единицу за ребро $v \leftrightarrow c$». Тогда считаются пары длины $k + 1$, не $k$.
3. Слияние без сдвига: $\mathrm{dp}[v][d] \mathrel{+}= \mathrm{dp}[c][d]$ вместо $\mathrm{dp}[v][d + 1] \mathrel{+}= \mathrm{dp}[c][d]$. Тогда вершины из $T_c$ числятся «на том же расстоянии от $v$, что и от $c$» — неверно, нужно $+1$ за ребро.
4. Проход по соседям без пропуска предка. Дерево хранится как неориентированный граф (списки смежности), и в for c in adj[v] оказывается непосредственный предок $v$. Без проверки if c == parent: continue — бесконечная рекурсия. Решение — либо передавать parent параметром, либо хранить ориентированное дерево.
5. int вместо long long в C++. Для $n = 50\,000$ ответ может достигать $\binom{n}{2} \approx 1.25 \cdot 10^9$ и более при больших $k$ на специфичных конфигурациях. Промежуточное произведение в свёртке может временно быть ещё больше. int32 достаточно близок к $2^{31}$, чтобы рискованно — берите long long.
6. Не освобождать dp[c] после слияния. На пределе памяти ($256$ МБ) полное хранение $n$ массивов длины $k + 1$ выливается в $200$ МБ; в Python с overhead'ом — гораздо больше.
7. Рекурсия глубины $50\,000$ в Python. Дефолтный лимит — $1000$. На цепочке падает почти сразу с RecursionError. Поднятие лимита через sys.setrecursionlimit помогает только до пределов стека ОС; итеративный DFS надёжнее.
8. Считать «упорядоченные пары» вместо «неупорядоченных». Если кажется, что свёртка симметрична и считает каждую пару дважды — это иллюзия. На самом деле каждая пара учитывается ровно один раз: при сшивании поддерева, содержащего «второй» из её концов в порядке обхода. Никакого дополнительного деления на 2 не нужно.

---

Анализ сложности

• Время: $O(n \cdot k)$. Каждое ребро $(v, c)$ обрабатывается ровно один раз, и работа в этот момент — два цикла по $k$ итераций (свёртка и слияние). Всего рёбер $n - 1$. Для $n = 50\,000$, $k = 500$ — $\sim 2.5 \cdot 10^7$ операций; на C++ укладывается в $1$ секунду с большим запасом.
• Память: в худшем случае (цепочка) — $O(n \cdot k)$ при наивном хранении всех $\mathrm{dp}[v]$, и $O(\text{глубина дерева} \cdot k)$ при освобождении $\mathrm{dp}[c]$ после слияния. Для сбалансированного дерева — $O(\log n \cdot k)$, для цепочки — те же $O(n \cdot k)$.

---

Что ещё полезно потренировать

Задачи на ту же идею — DP на дереве с накоплением подответов через свёртку или max/sum-комбинацию. Все — с возраст-нейтральных платформ.

• Codeforces Div.2 D/E уровня $1800$–$2200$ с тегом dp on trees: классические постановки «найти максимум / посчитать число / выбрать подмножество вершин с ограничением». Шаблон тот же — DFS post-order, в каждой вершине агрегируются ответы детей.
• Codeforces задача «independent set» на дереве ($1500$–$1800$): состояние из двух чисел «вершина взята / не взята», переход — из ответов детей. Самая лёгкая для входа в технику.
• acmp.ru, раздел «Деревья»: задачи $1300$–$1700$ на простые tree DP — поиск диаметра, числа листьев, суммарного расстояния.
• Centroid decomposition — альтернативный подход к этой же задаче с асимптотикой $O(n \log n)$ независимо от $k$. Сложнее в реализации, но окупается на задачах с $k \le n$ (когда $n \cdot k$ становится квадратичным).

Для следующего шага — попробовать tree DP, где состояние двумерное (например, «количество вершин, выбранных в подмножество, $\le j$») или где помимо числа вершин нужно агрегировать сумму, минимум, дисперсию.

---

Итого

• Идея: подвешиваем дерево за корень, в каждой вершине $v$ держим массив $\mathrm{dp}[v][d]$ — число вершин в поддереве на расстоянии $d$. При сшивании очередной дочерней вершины $c$ сначала добавляем к ответу свёртку $\sum_{i + j = k - 1} \mathrm{dp}[v][i] \cdot \mathrm{dp}[c][j]$ (это пары длины $k$ с LCA в $v$), потом сливаем $\mathrm{dp}[c]$ в $\mathrm{dp}[v]$ со сдвигом $+1$.
• Сложность: $O(n \cdot k)$ по времени, $O(\text{глубина дерева} \cdot k)$ по памяти при освобождении $\mathrm{dp}[c]$ после слияния.
• Ловушки: перепутанный порядок «свёртка → слияние», свёртка $i + j = k$ вместо $k - 1$, слияние без сдвига, рекурсия в Python на цепочке, int вместо long long в C++.
• Стороннее замечание: для $k \sim n$ алгоритм $O(n k)$ деградирует до $O(n^2)$ — здесь нужен переход на centroid decomposition.