Heavy-light декомпозиция: запросы на пути дерева за O(log² n)

Коротко

Задача

Дано дерево, у вершин есть значения. Нужно уметь:

• узнать сумму (или максимум) на пути между двумя вершинами;
• изменить значение вершины.

Наивно путь проходится за O(n) — при 10⁵ запросах это безнадёжно.

Heavy-light декомпозиция сводит задачу к отрезкам массива, где уже работает дерево отрезков.

Идея: тяжёлые и лёгкие рёбра

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

            1
          / | \
        2   3   4        размеры поддеревьев: 2→5, 3→1, 4→1
       / \
      5   6              тяжёлое ребро: 1—2, затем 2—5
     /
    7

Тяжёлые рёбра образуют непересекающиеся вертикальные цепи, покрывающие всё дерево.

Ключевое свойство: любой путь от вершины до корня пересекает не более log₂n лёгких рёбер.

Доказательство простое: переход по лёгкому ребру означает спуск в поддерево, размер которого меньше половины родительского (иначе ребро было бы тяжёлым). Размер не может делиться пополам больше log₂n раз.

Отсюда: путь проходит через O(log n) цепей, а внутри каждой цепи это непрерывный отрезок массива — запрос к дереву отрезков за O(log n). Итого O(log² n).

Нумерация вершин

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

import sys


class HLD:
    def __init__(self, tree: dict, root: int, values: dict):
        self.tree = tree
        self.parent = {root: -1}
        self.depth = {root: 0}
        self.heavy = {}          # вершина -> тяжёлый ребёнок
        self.head = {}           # вершина -> начало её цепи
        self.pos = {}            # вершина -> позиция в массиве
        self.n = len(tree)

        self._compute_sizes(root)
        self._decompose(root, root)

        # массив значений в порядке нумерации
        self.base = [0] * self.n
        for v, p in self.pos.items():
            self.base[p] = values[v]

        self.seg = SegmentTree(self.base)

    def _compute_sizes(self, root) -> None:
        # итеративный обход: сначала порядок, потом размеры снизу вверх
        order, stack = [], [root]
        while stack:
            v = stack.pop()
            order.append(v)
            for u in self.tree[v]:
                if u != self.parent[v]:
                    self.parent[u] = v
                    self.depth[u] = self.depth[v] + 1
                    stack.append(u)

        size = {v: 1 for v in order}
        for v in reversed(order):
            best, best_size = -1, 0
            for u in self.tree[v]:
                if u == self.parent[v]:
                    continue
                size[v] += size[u]
                if size[u] > best_size:
                    best, best_size = u, size[u]
            if best != -1:
                self.heavy[v] = best

        self.order = order

    def _decompose(self, root, chain_head) -> None:
        counter = 0
        stack = [(root, chain_head)]

        while stack:
            v, h = stack.pop()
            # идём вниз по тяжёлым рёбрам — вся цепь получает подряд номера
            while True:
                self.head[v] = h
                self.pos[v] = counter
                counter += 1

                # лёгкие дети начинают новые цепи
                for u in self.tree[v]:
                    if u != self.parent[v] and u != self.heavy.get(v):
                        stack.append((u, u))

                if v in self.heavy:
                    v = self.heavy[v]
                else:
                    break

    def query_path(self, u: int, v: int) -> int:
        """Сумма на пути между u и v."""
        result = 0

        while self.head[u] != self.head[v]:
            # поднимаем ту вершину, чья цепь начинается глубже
            if self.depth[self.head[u]] < self.depth[self.head[v]]:
                u, v = v, u
            result += self.seg.query(self.pos[self.head[u]], self.pos[u])
            u = self.parent[self.head[u]]

        # обе в одной цепи — один отрезок
        if self.depth[u] > self.depth[v]:
            u, v = v, u
        result += self.seg.query(self.pos[u], self.pos[v])

        return result

    def update(self, v: int, value: int) -> None:
        self.seg.update(self.pos[v], value)


class SegmentTree:
    def __init__(self, arr: list[int]):
        self.n = len(arr)
        self.tree = [0] * (2 * self.n)
        for i, x in enumerate(arr):
            self.tree[self.n + i] = x
        for i in range(self.n - 1, 0, -1):
            self.tree[i] = self.tree[2 * i] + self.tree[2 * i + 1]

    def update(self, i: int, value: int) -> None:
        i += self.n
        self.tree[i] = value
        i //= 2
        while i:
            self.tree[i] = self.tree[2 * i] + self.tree[2 * i + 1]
            i //= 2

    def query(self, l: int, r: int) -> int:
        """Сумма на [l, r] включительно."""
        result = 0
        l += self.n
        r += self.n + 1
        while l < r:
            if l & 1:
                result += self.tree[l]
                l += 1
            if r & 1:
                r -= 1
                result += self.tree[r]
            l //= 2
            r //= 2
        return result


tree = {
    1: [2, 3, 4],
    2: [1, 5, 6],
    3: [1],
    4: [1],
    5: [2, 7],
    6: [2],
    7: [5],
}
values = {1: 10, 2: 20, 3: 30, 4: 40, 5: 50, 6: 60, 7: 70}

hld = HLD(tree, root=1, values=values)

print(hld.query_path(7, 3))   # 10+20+50+70+30 = 180
print(hld.query_path(6, 4))   # 60+20+10+40 = 130

hld.update(2, 200)
print(hld.query_path(7, 3))   # 360

Разбор запроса

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

Условие «поднимаем ту, чья голова цепи глубже» гарантирует, что мы не проскочим точку слияния — фактически мы движемся к наименьшему общему предку.

Что можно считать

Помимо суммы подходит любая ассоциативная операция: максимум, минимум, НОД, XOR. Достаточно поменять операцию в дереве отрезков.

Отдельный случай — обновление на пути (прибавить ко всем вершинам пути). Тогда нужно дерево отрезков с отложенными операциями, сложность запроса остаётся O(log² n).

Альтернативы

HLD — практичный выбор: логарифм в квадрате обычно достаточно быстр, а код в разы проще link-cut деревьев.

Если дерево меняется (рёбра добавляются и удаляются), HLD не подойдёт — нужны link-cut деревья.

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

Выбор тяжёлого ребёнка не по размеру поддерева. Свойство «не более log n лёгких рёбер» держится только на размерах. Выбор по глубине или произвольный сломает оценку.

Нумерация без учёта цепей. Если вершины одной цепи не получат подряд идущие номера, отрезок в массиве не сложится.

Подъём не той вершины. Сравнивать нужно глубину головы цепи, а не самой вершины.

Рекурсивный обход на глубоких деревьях. В коде намеренно итеративный.

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

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

• HLD разбивает дерево на вертикальные цепи по тяжёлым рёбрам.
• Тяжёлое ребро ведёт к ребёнку с наибольшим поддеревом.
• Путь пересекает не более log n лёгких рёбер — отсюда O(log² n) на запрос.
• Вершины одной цепи нумеруются подряд, что превращает цепь в отрезок массива.
• Поверх массива работает обычное дерево отрезков.

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

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

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

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

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