Центроидная декомпозиция: задачи о путях в дереве за O(n log n)

Коротко

Что такое центроид

Центроид дерева — вершина, при удалении которой все получившиеся части имеют размер не более n/2.

        1
      / | \
     2  3  4
    /
   5

центроид — вершина 1: части размером 2, 1, 1

Ключевые факты:

• центроид существует всегда;
• их может быть не больше двух;
• находится за O(n) одним обходом.

Идея декомпозиции

Находим центроид, «удаляем» его, рекурсивно обрабатываем получившиеся части. Центроиды образуют новое дерево — дерево центроидов.

Поскольку каждое удаление делит дерево минимум пополам, глубина дерева центроидов равна O(log n). Это и есть источник логарифма.

Главное свойство, ради которого всё делается:

Любой путь в исходном дереве проходит через центроид одного из уровней декомпозиции.

Значит, обработав для каждого центроида все пути, которые через него проходят, мы покроем все пути дерева. А суммарный размер обрабатываемых частей на каждом уровне — O(n), уровней O(log n), итого O(n log n).

Поиск центроида

def find_centroid(tree: dict, start: int, removed: set) -> int:
    # считаем размеры поддеревьев
    size = {}
    order, parent = [], {start: -1}
    stack = [start]

    while stack:
        v = stack.pop()
        order.append(v)
        for u in tree[v]:
            if u != parent[v] and u not in removed:
                parent[u] = v
                stack.append(u)

    for v in reversed(order):
        size[v] = 1 + sum(size[u] for u in tree[v]
                          if u != parent[v] and u not in removed)

    total = size[start]

    # ищем вершину, у которой все части <= total/2
    v = start
    while True:
        heavy = -1
        for u in tree[v]:
            if u != parent[v] and u not in removed and size[u] > total // 2:
                heavy = u
                break
        if heavy == -1:
            return v
        parent[heavy] = v
        v = heavy

Спуск в «тяжёлого» ребёнка — стандартный приём: центроид всегда лежит в направлении наибольшего поддерева.

Построение дерева центроидов

def build_centroid_tree(tree: dict, root: int):
    removed = set()
    centroid_parent = {}

    def decompose(entry: int, parent_centroid: int) -> None:
        c = find_centroid(tree, entry, removed)
        centroid_parent[c] = parent_centroid
        removed.add(c)

        for u in tree[c]:
            if u not in removed:
                decompose(u, c)

    decompose(root, -1)
    return centroid_parent


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

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

Вершина 2 стала корнем дерева центроидов, хотя в исходном дереве корнем была единица. Это нормально — структуры независимы.

Классическая задача: пути заданной длины

Посчитать количество пар вершин, расстояние между которыми ровно k.

Наивно — O(n²). С центроидной декомпозицией — O(n log n).

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

from collections import defaultdict


def count_paths_of_length(tree: dict, root: int, k: int) -> int:
    removed = set()
    total = 0

    def collect_depths(start: int, banned: int) -> list[int]:
        """Глубины всех вершин части, начиная со start."""
        depths = []
        stack = [(start, banned, 1)]
        while stack:
            v, p, d = stack.pop()
            depths.append(d)
            for u in tree[v]:
                if u != p and u not in removed:
                    stack.append((u, v, d + 1))
        return depths

    def count_pairs(depths: list[int]) -> int:
        counter = defaultdict(int)
        for d in depths:
            counter[d] += 1
        pairs = 0
        for d, cnt in counter.items():
            other = k - d
            if other == d:
                pairs += cnt * (cnt - 1) // 2
            elif other in counter and other > d:
                pairs += cnt * counter[other]
        return pairs

    def decompose(entry: int) -> None:
        nonlocal total
        c = find_centroid(tree, entry, removed)
        removed.add(c)

        all_depths = []
        for u in tree[c]:
            if u in removed:
                continue
            sub = collect_depths(u, c)
            # вычитаем пары внутри одного поддерева
            total -= count_pairs(sub)
            all_depths.extend(sub)
            # пути, у которых один конец — сам центроид
            total += sum(1 for d in sub if d == k)

        total += count_pairs(all_depths)

        for u in tree[c]:
            if u not in removed:
                decompose(u)

    decompose(root)
    return total


path_tree = {1: [2], 2: [1, 3], 3: [2, 4], 4: [3]}
print(count_paths_of_length(path_tree, 1, 2))   # 2 → (1,3) и (2,4)

Приём «посчитать все пары, вычесть пары внутри одного поддерева» — стандартный для центроидной декомпозиции. Он встречается почти во всех задачах на эту тему.

Где применяется

Подсчёт путей заданной длины, с суммой весов, с ограничением.

Ближайшая отмеченная вершина. Для каждого запроса подниматься по дереву центроидов и хранить минимумы — O(log n) на запрос.

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

Разбиение задачи «разделяй и властвуй» на деревьях — общий каркас.

Сравнение с HLD

Это разные инструменты. HLD отвечает «какая сумма на пути от u до v». Центроидная декомпозиция отвечает «сколько всего путей с таким-то свойством».

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

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

Забытое вычитание пар из одного поддерева. Даёт завышенный ответ — пути, не проходящие через центроид, засчитываются дважды.

Поиск центроида за O(n) на каждом уровне без ограничения частью. Суммарно даст O(n²) вместо O(n log n).

Рекурсия глубиной n. Дерево центроидов имеет глубину log n, но collect_depths по исходному дереву может уйти глубоко — используйте итеративный обход.

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

• Центроид делит дерево на части размером не более половины.
• Дерево центроидов имеет глубину O(log n).
• Любой путь исходного дерева проходит через центроид какого-то уровня.
• Типичный приём: посчитать все пары через центроид, вычесть пары внутри одного поддерева.
• Решает задачи о множестве путей, в отличие от HLD, который работает с конкретным путём.

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

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

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

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

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