Дерево отрезков

Дерево отрезков (англ. segment tree) — мощная и гибкая структура данных, позволяющая быстро отвечать на разнообразные запросы на отрезках массива и менять его элементы.

В качестве разминочной задачи рассмотрим следующую: дан массив \(a\) из \(n\) целых чисел, и требуется отвечать на запросы двух типов:

  1. Изменить значение в ячейке: a[k] = x.
  2. Посчитать сумму элементов на полуинтервале \([l, r)\).

Оба запроса нужно обрабатывать за время \(O(\log n)\).

Структура дерева отрезков

Чтобы решить задачу, проделаем с массивом следующие манипуляции.

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

Эту последовательность разбиений можно представить в виде дерева:

Корень соответствует отрезку \([0, n)\), а каждая внутренняя вершина имеет ровно двух сыновей, которые тоже соответствуют каким-то отрезкам — половинам родительского. Отсюда и название: «дерево отрезков».

Свойства дерева отрезков

Высота дерева отрезков равна \(\Theta(\log n)\): на каждом уровне длина отрезка уменьшается вдвое. Этот факт ключевой для оценки асимптотики операций.

Память. Дерево содержит менее \(2n\) вершин: первый уровень содержит одну вершину (корень), второй — не более двух, третий — не более четырёх, и так далее. В худшем случае количество вершин не превосходит \(n + n/2 + n/4 + \ldots + 1 < 2n\). На практике обычно выделяют массив размера \(4n\), чтобы гарантированно хватило в случае, когда \(n\) — не степень двойки.

Разбиение отрезка. Любой полуинтервал \([l, r)\) разбивается на \(O(\log n)\) непересекающихся отрезков, соответствующих вершинам дерева: с каждого уровня нам достаточно не более двух отрезков.

При \(n\), отличных от степеней двойки, не все уровни дерева будут полностью заполнены. Например, при \(n = 3\) левый сын корня — отрезок \([0, 2)\) — имеет двух потомков, в то время как правый сын — отрезок \([2, 3)\) — является листом.

Алгоритмы операций

Запрос обновления. Нужно обновить значения в вершинах так, чтобы они соответствовали новому значению a[k] = x. Изменятся только те вершины, в чей отрезок входит позиция \(k\) — по одной с каждого уровня, итого \(\Theta(\log n)\) штук.

Это удобно реализовать рекурсивной функцией: ей передаётся текущая вершина дерева, она рекурсивно вызывает себя от того сына, в чей отрезок входит \(k\), а после возврата пересчитывает свою сумму как сумму значений в детях.

Запрос суммы. Считаем, что во всех вершинах хранятся корректные суммы. Тоже воспользуемся рекурсивной функцией, рассмотрев три случая:

  1. Отрезок вершины целиком лежит в отрезке запроса — вернуть сумму, записанную в вершине.
  2. Отрезки вершины и запроса не пересекаются — вернуть 0.
  3. Иначе рекурсивно вызваться от обоих сыновей и сложить их ответы.

Чтобы понять, почему это работает за \(O(\log n)\), нужно оценить количество «интересных» отрезков — тех, которые порождают новые вызовы рекурсии. Это в точности отрезки, содержащие границу запроса: остальные либо целиком лежат внутри, либо целиком снаружи и сразу завершаются. Каждая из двух границ запроса содержится в \(O(\log n)\) отрезках, поэтому итоговая асимптотика — \(O(\log n)\).

Реализация

Соглашения. Мы будем использовать полуинтервалы \([l, r)\) вместо отрезков. Несмотря на контринтуитивность, это упрощает код и в целом является хорошей практикой, как и нумерация с нуля.

Хранение в массиве. Так же, как и в куче, дерево отрезков удобно хранить в массиве tree. Корень — это tree[1], а у вершины с индексом \(v\) дети находятся по индексам \(2v\) и \(2v + 1\). Размера \(4n\) всегда достаточно.

class SegmentTree {
public:
    SegmentTree(const std::vector<int64_t>& a) : n(a.size()), tree(4 * n) {
        Build(1, 0, n, a);
    }

    void Update(size_t pos, int64_t value) {
        Update(1, 0, n, pos, value);
    }

    int64_t Sum(size_t l, size_t r) {
        return Sum(1, 0, n, l, r);
    }

private:
    size_t n;
    std::vector<int64_t> tree;

    void Build(size_t v, size_t l, size_t r, const std::vector<int64_t>& a) {
        if (r - l == 1) {
            tree[v] = a[l];
            return;
        }
        size_t mid = (l + r) / 2;
        Build(2 * v, l, mid, a);
        Build(2 * v + 1, mid, r, a);
        tree[v] = tree[2 * v] + tree[2 * v + 1];
    }

    void Update(size_t v, size_t l, size_t r, size_t pos, int64_t value) {
        if (r - l == 1) {
            tree[v] = value;
            return;
        }
        size_t mid = (l + r) / 2;
        if (pos < mid) {
            Update(2 * v, l, mid, pos, value);
        } else {
            Update(2 * v + 1, mid, r, pos, value);
        }
        tree[v] = tree[2 * v] + tree[2 * v + 1];
    }

    int64_t Sum(size_t v, size_t l, size_t r, size_t ql, size_t qr) {
        if (qr <= l || r <= ql) {
            return 0;
        }
        if (ql <= l && r <= qr) {
            return tree[v];
        }
        size_t mid = (l + r) / 2;
        return Sum(2 * v, l, mid, ql, qr) + Sum(2 * v + 1, mid, r, ql, qr);
    }
};

Заметим, что публичный интерфейс класса не требует от пользователя знать о внутреннем устройстве дерева — про индекс вершины и её границы знают только приватные функции.

Сложность операций

Операция Сложность
Построение \(O(n)\)
Точечное обновление \(O(\log n)\)
Сумма на отрезке \(O(\log n)\)
Память \(O(n)\)

Обобщение на другие операции

Дерево отрезков работает не только для суммы. Достаточно, чтобы операция была ассоциативной — то есть чтобы \((a \star b) \star c = a \star (b \star c)\). В этом случае значение в каждой вершине пересчитывается одинаково: \(\text{tree}[v] = \text{tree}[2v] \star \text{tree}[2v + 1]\).

Примеры подходящих операций:

  • сумма;
  • минимум и максимум;
  • НОД и НОК;
  • побитовые AND, OR, XOR;
  • произведение по модулю.

Для каждой такой операции дерево отрезков отвечает на запрос «применить операцию ко всем элементам отрезка» за \(O(\log n)\).

Оптимизации

Описанная реализация простая и расширяемая, но не самая быстрая. Несколько идей, как ускорить её:

  • Не хранить границы \(l\) и \(r\) в вершине, а пересчитывать во время спуска (мы так уже и сделали).
  • Перейти от рекурсии к итеративной реализации «снизу вверх», работающей без явного дерева.

Подробнее об этих и других оптимизациях можно почитать у Е-макса и на CodeForces.

Отложенные изменения

Часто хочется не только менять отдельные элементы, но и применять одну и ту же операцию сразу ко всему отрезку — например, «присвоить всем элементам \([l, r)\) значение \(x\)». Если делать это «в лоб», спускаясь до листьев, в худшем случае мы изменим \(O(n)\) вершин — теряется вся выгода от структуры.

Решение — техника отложенных изменений (англ. lazy propagation).

Идея

Вместо того чтобы сразу менять все вершины, будем «помечать» поддерево пометкой \(\text{lazy}_v = x\), означающей: «всем элементам этого отрезка должно быть присвоено \(x\), но я ещё не успел этого сделать».

Например, если пришёл запрос «присвоить \(x\) на всём массиве», мы вообще ничего не изменяем — только ставим пометку в корне.

Когда позже понадобятся правильные значения детей, мы по пути от корня будем выполнять «проталкивание» (push): пересчитаем сумму текущего отрезка по пометке и передадим её детям. Спускаться рекурсивно «впрок» при этом не нужно — проталкивание сработает у детей в нужный момент.

Экономия получается за счёт того, что при передаче новой пометки в ребёнка старая просто перезаписывается: некоторые отложенные операции в итоге так и не понадобится выполнять.

Асимптотика остаётся той же \(O(\log n)\) на запрос: мы просто делаем константно больше работы в каждой вершине на пути.

Реализация

Добавим к каждой вершине поле lazy для отложенного присваивания. Будем считать, что значение lazy = -1 означает «пометки нет» (это работает, если допустимые значения элементов неотрицательны — иначе нужен отдельный флаг или std::optional).

class SegmentTree {
public:
    SegmentTree(const std::vector<int64_t>& a)
        : n(a.size()), tree(4 * n), lazy(4 * n, -1) {
        Build(1, 0, n, a);
    }

    void Assign(size_t l, size_t r, int64_t value) {
        Assign(1, 0, n, l, r, value);
    }

    int64_t Sum(size_t l, size_t r) {
        return Sum(1, 0, n, l, r);
    }

private:
    size_t n;
    std::vector<int64_t> tree;
    std::vector<int64_t> lazy;

    void Build(size_t v, size_t l, size_t r, const std::vector<int64_t>& a) {
        if (r - l == 1) {
            tree[v] = a[l];
            return;
        }
        size_t mid = (l + r) / 2;
        Build(2 * v, l, mid, a);
        Build(2 * v + 1, mid, r, a);
        tree[v] = tree[2 * v] + tree[2 * v + 1];
    }

    void Push(size_t v, size_t l, size_t r) {
        if (lazy[v] == -1) return;
        tree[v] = (r - l) * lazy[v];
        if (r - l > 1) {
            lazy[2 * v] = lazy[v];
            lazy[2 * v + 1] = lazy[v];
        }
        lazy[v] = -1;
    }

    void Assign(size_t v, size_t l, size_t r, size_t ql, size_t qr, int64_t value) {
        Push(v, l, r);
        if (qr <= l || r <= ql) return;
        if (ql <= l && r <= qr) {
            lazy[v] = value;
            Push(v, l, r);
            return;
        }
        size_t mid = (l + r) / 2;
        Assign(2 * v, l, mid, ql, qr, value);
        Assign(2 * v + 1, mid, r, ql, qr, value);
        tree[v] = tree[2 * v] + tree[2 * v + 1];
    }

    int64_t Sum(size_t v, size_t l, size_t r, size_t ql, size_t qr) {
        Push(v, l, r);
        if (qr <= l || r <= ql) return 0;
        if (ql <= l && r <= qr) return tree[v];
        size_t mid = (l + r) / 2;
        return Sum(2 * v, l, mid, ql, qr) + Sum(2 * v + 1, mid, r, ql, qr);
    }
};

Ключевой инвариант: вызвав Push в самом начале обработки вершины, мы гарантируем, что в ней лежит корректное значение tree[v], а в её детях достаточно информации, чтобы при дальнейшем спуске и вызове Push от них значения тоже стали корректными.

По-английски эта техника называется lazy propagation. Общая идея «делать всё в последний момент» применима не только в дереве отрезков, но и в других структурах данных.