Динамическое программирование на деревьях: разбор с примерами

Коротко

Главное свойство, на котором всё держится: у дерева нет циклов, поэтому каждая вершина обрабатывается ровно один раз, и подзадачи не пересекаются.

Базовая схема

ДП на дереве почти всегда выглядит одинаково: обход в глубину, вычисление ответа для вершины на основе ответов детей.

def dfs(v, parent):
    for child in graph[v]:
        if child != parent:
            dfs(child, v)
            # объединяем ответ ребёнка с ответом вершины
    # финализируем dp[v]

Проверка child != parent заменяет массив посещённых — в дереве вернуться можно только туда, откуда пришёл.

Пример 1: размеры поддеревьев

Самая простая задача, которая нужна как заготовка для многих других.

import sys
from collections import defaultdict


def subtree_sizes(graph: dict, root: int) -> dict:
    size = {}

    # итеративный обход, чтобы не упереться в стек
    order = []
    parent = {root: None}
    stack = [root]

    while stack:
        v = stack.pop()
        order.append(v)
        for u in graph[v]:
            if u != parent[v]:
                parent[u] = v
                stack.append(u)

    # обрабатываем в обратном порядке: дети раньше родителей
    for v in reversed(order):
        size[v] = 1 + sum(size[u] for u in graph[v] if u != parent[v])

    return size


tree = {
    1: [2, 3],
    2: [1, 4, 5],
    3: [1],
    4: [2],
    5: [2],
}

print(subtree_sizes(tree, 1))   # {4: 1, 5: 1, 3: 1, 2: 3, 1: 5}

Приём с обратным порядком обхода очень удобен: он даёт гарантию, что дети посчитаны раньше родителя, без рекурсии.

Пример 2: диаметр дерева

Диаметр — длина самого длинного пути между двумя вершинами.

Ключевая мысль: любой путь в дереве имеет самую высокую точку — вершину, через которую он проходит. Для каждой вершины посчитаем две самые длинные «ветки вниз» и сложим их.

def tree_diameter(graph: dict, root: int) -> int:
    depth = {}       # максимальная глубина вниз от вершины
    diameter = 0

    order, parent = [], {root: None}
    stack = [root]
    while stack:
        v = stack.pop()
        order.append(v)
        for u in graph[v]:
            if u != parent[v]:
                parent[u] = v
                stack.append(u)

    for v in reversed(order):
        # две самые длинные ветки вниз
        best_two = sorted(
            (depth[u] for u in graph[v] if u != parent[v]),
            reverse=True
        )[:2]

        depth[v] = 1 + (best_two[0] if best_two else -1)
        # путь через v — сумма двух лучших веток
        nonlocal_path = sum(b + 1 for b in best_two)
        diameter = max(diameter, nonlocal_path)

    return diameter


print(tree_diameter(tree, 1))   # 3  → путь 4-2-1-3


bamboo = {1: [2], 2: [1, 3], 3: [2, 4], 4: [3]}
print(tree_diameter(bamboo, 1))   # 3

Есть и более простой способ найти диаметр — два обхода в ширину: из произвольной вершины найти самую дальнюю, из неё найти самую дальнюю снова. Расстояние между ними и есть диаметр. Доказательство нетривиально, но код совсем короткий.

Пример 3: максимальное независимое множество

Классическая задача: выбрать максимальное число вершин так, чтобы никакие две не были соединены ребром. Формулировка «вечеринка в компании»: пригласить как можно больше сотрудников, но никто не должен прийти вместе со своим непосредственным начальником.

Здесь у каждой вершины два состояния:

• dp[v][0] — максимум в поддереве v, если саму v не берём;
• dp[v][1] — максимум, если v берём.

Переходы:

dp[v][1] = 1 + Σ dp[child][0]     взяли v → детей брать нельзя
dp[v][0] =     Σ max(dp[child][0], dp[child][1])   не взяли v → дети свободны

def max_independent_set(graph: dict, root: int) -> int:
    take = {}      # dp[v][1]
    skip = {}      # dp[v][0]

    order, parent = [], {root: None}
    stack = [root]
    while stack:
        v = stack.pop()
        order.append(v)
        for u in graph[v]:
            if u != parent[v]:
                parent[u] = v
                stack.append(u)

    for v in reversed(order):
        children = [u for u in graph[v] if u != parent[v]]
        take[v] = 1 + sum(skip[u] for u in children)
        skip[v] = sum(max(skip[u], take[u]) for u in children)

    return max(take[root], skip[root])


print(max_independent_set(tree, 1))   # 3  → вершины 3, 4, 5

Схема «два состояния на вершину» покрывает огромный класс задач: выбор вершин с ограничениями, покрытие, доминирующее множество.

Переподвешивание (rerooting)

Более продвинутый приём. Иногда нужен ответ для каждой вершины как корня — например, «для каждой вершины найти расстояние до самой дальней».

Наивно — запустить обход из каждой вершины, O(n²). Техника переподвешивания даёт O(n).

Идея в два прохода:

1. Снизу вверх считаем ответы для поддеревьев при фиксированном корне.
2. Сверху вниз пересчитываем ответы, «переворачивая» ребро: для ребёнка добавляем вклад «всего остального дерева», который вычисляется из ответа родителя.

def farthest_distance_from_each(graph: dict, root: int) -> dict:
    down = {}      # максимальное расстояние вниз
    up = {}        # максимальное расстояние вверх (через родителя)

    order, parent = [], {root: None}
    stack = [root]
    while stack:
        v = stack.pop()
        order.append(v)
        for u in graph[v]:
            if u != parent[v]:
                parent[u] = v
                stack.append(u)

    # проход 1: снизу вверх
    for v in reversed(order):
        children = [u for u in graph[v] if u != parent[v]]
        down[v] = max((down[u] + 1 for u in children), default=0)

    # проход 2: сверху вниз
    up[root] = 0
    for v in order:
        children = [u for u in graph[v] if u != parent[v]]

        # два лучших спуска из v, чтобы исключить вклад самого ребёнка
        best = sorted((down[u] + 1 for u in children), reverse=True)[:2]

        for u in children:
            own = down[u] + 1
            # лучший спуск из v, не через u
            other = best[0] if best[0] != own else (best[1] if len(best) > 1 else 0)
            up[u] = max(up[v], other) + 1

    return {v: max(down[v], up[v]) for v in graph}


print(farthest_distance_from_each(tree, 1))
# {1: 2, 2: 2, 3: 3, 4: 3, 5: 3}

Тонкий момент — при вычислении up[u] нужно исключить вклад самого u, иначе получится путь, идущий вниз и сразу возвращающийся. Отсюда трюк с двумя лучшими значениями.

Почему рекурсию лучше заменить на итерацию

Естественная реализация ДП на дереве рекурсивна. Но на дереве-«бамбуке» из 10⁵ вершин глубина рекурсии равна 10⁵, и Python упадёт с RecursionError при лимите по умолчанию в 1000.

import sys
sys.setrecursionlimit(300000)

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

Типичные задачи

• Размер и сумма поддерева — база для многих других.
• Диаметр и центр дерева.
• Максимальное независимое множество, минимальное вершинное покрытие.
• Раскраска дерева в k цветов.
• Количество путей заданной длины.
• Наибольшая сумма пути между любыми двумя вершинами.
• Задача о ранце на дереве — выбрать поддерево с ограничением по весу.

Частые ошибки

Отсутствие проверки на родителя. Без if child != parent обход зациклится между двумя соседними вершинами.

Обработка родителя раньше детей. Порядок принципиален: ответ вершины строится из ответов детей.

Забытый случай листа. У листа нет детей — max() по пустой последовательности упадёт. Используйте default=.

Некорректное исключение вклада ребёнка при переподвешивании. Самая тонкая ошибка: без исключения путь «схлопывается» сам в себя, и ответ завышается.

Рекурсия на глубоких деревьях. Разобрано выше.

Что запомнить

• ДП на дереве — обход в глубину с вычислением ответа вершины из ответов детей.
• Проверка child != parent заменяет массив посещённых.
• Схема «два состояния на вершину» решает задачи выбора с ограничениями.
• Переподвешивание даёт ответ для всех корней за O(n) вместо O(n²).
• На больших деревьях используйте итеративный обход с обработкой в обратном порядке.

Пройди собеседование в топ-компанию

Платформа для подготовки

Решай алгоритмические задачи как профи

✓ Популярные алгоритмы✓ Разбор решений✓ AI помощь

Начать сейчас→