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

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

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

Рассмотрим конкретную задачу:

Дан массив \(a\) из \(n\) целых чисел, нужно уметь отвечать на запросы двух типов:

1. Изменить значение в ячейке (т. е. отреагировать на присвоение a[k] = x).

2. Вывести сумму элементов \(a_i\) на отрезке с \(l\) по \(r\).

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

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

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

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

Строить его можно рекурсивной функцией:

• Если вершина является листом, взять в качестве суммы значение соответствующей ячейки.
• Если вершина является отрезком, разделить его на два и в качестве суммы взять сумму его детей.

Разные свойства

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

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

Дерево также содержит менее \(2n\) вершин: первый уровень дерева отрезков содержит одну вершину (корень), второй уровень — в худшем случае две вершины, на третьем уровне в худшем случае будет четыре вершины, и так далее, пока число вершин не достигнет \(n\). Таким образом, число вершин в худшем случае оценивается суммой \(n + \frac{n}{2} + \frac{n}{4} + \frac{n}{8} + \ldots + 1 < 2n\). Значит, оно линейное по памяти.

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

Ок, как это нам поможет?

Опишем теперь, как с помощью такой структуры решить задачу.

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

Изменим все вершины, в суммах которых участвует \(k\)-тый элемент. Их будет \(\Theta(\log n)\) — по одной с каждого уровня.

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

Запрос суммы. Мы знаем, что во всех вершинах лежат корректные значения.

Сделаем тоже рекурсивную функцию, рассмотрев три случая:

• Если отрезок вершины лежит целиком в отрезке запроса, то вернуть записанную в ней сумму.
• Если отрезки вершины и запроса не пересекаются, то вернуть 0.
• Иначе разделиться рекурсивно на 2 и вернуть сумму этой функции от обоих детей.

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

Ликбез по C++

Наша реализация будет на указателях. Никто не говорит, что она самая лучшая (см. раздел «Другие реализации»), но она самая понятная. Вам может поначалу показаться, что она слишком сложная, но позже вы поймёте её преимущества.

Но сначала нам нужно рассказать про объектно-ориентированное программирование и некоторые фишки C++. Если вы их уже знаете, то можете пропускать этот раздел.

Объект — это сущность, которой можно посылать сообщения и которая может на них реагировать, используя свои данные. Инкапсулировать логику в объекты на самом деле очень удобно. Дереву отрезков не важно знать, как устроен окружающий мир, а миру не важно, как внутри устроено дерево отрезков — это просто какая-то структура, которая умеет делать нужные операции за \(O(\log n)\).

В C++ есть два способа объявлять классы (объект — это экземпляр класса): через struct и через class. Их основное отличие в том, что по умолчанию в class все поля приватные: к ним нет прямого доступа снаружи. Это нужно для дополнительной защиты, чтобы в крупных промышленных проектах никто случайно ничего не поломал, но на олимпиадах это не очень актуально.

У классов есть поля (переменные) и методы (функции, привязанные к объектам). Среди них есть особые, например конструктор — он вызывается при создании объекта. Чтобы объявить конструктор класса в C++, нужно объявить внутри него метод с тем же названием, что и у самого класса.

struct A {
    int param1, param2; // тут можно что-то хранить
    char param3 = 'k';
    A (int var) {
        // эта часть называется конструктором
        // ...
    }
    void do_something () {
        // это какой-то другой метод
        // ...
    }
}; // <- не забудьте точку с запятой

Другое важное понятие — указатель. Память можно представлять как просто очень большой массив. На самом деле, когда мы создаем какой-то объект, отдельная программа (аллокатор) выделяет место в массиве (оперативной памяти) под этот объект и возвращает позицию (указатель) на место в этом массиве.

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

A x(179);
x.do_something();
x.param1 = 57;

A *y = new A(42); // new возвращает адрес, по которому можно найти объект
y->do_something();
y.param3 = '!';

Оффтоп: вы не задумывались, почему мы перешли с 32-битных процессоров на 64-битные? Каждый указатель ссылается на байт — более точный адрес менеджер памяти выделять не умеет. Поэтому 32-битный компьютер умеет работать только с не более, чем \(2^{32}\) байтами памяти — это ровно 4 гигабайта — что с какого-то момента начало нехватать. Большинство операций в любом компьютере — это операции с памятью, и разрядность повысили именно из-за этого, а не чтобы операции с long long быстрее считались

Реализация

Примечание. Данная реализация неэффективна по времени и памяти (см. раздел «другие реализации»). Мы её приводим в педагогических целях.

Общий план реализации любых структур данных:

1. Полностью понять все инварианты — как должна выглядеть структура, какие значения должны принимать поля, etc.
2. Формально описать, что должны делать методы и за какую асимптотику.
3. Решить много отдельных задач, реализуя методы, не нарушающие инварианты.

struct segtree {
    int lb, rb; // левые и правые границы отрезков
    int sum = 0; // сумма на текущем отрезке
    segtree *l = 0, *r = 0;
    segtree (int _lb, int _rb) {
        lb = _lb, rb = _rb;
        if (lb + 1 < rb) {
            // если не лист, создаем детей
            int t = (lb + rb) / 2;
            l = new segtree(lb, t);
            r = new segtree(t, rb);
        }
    }
    void add (int k, int x) {
        sum += x;
        if (l) {
            if (k < l->rb)
                l->add(k, x);
            else
                r->add(k, x);
        }
    }
    int get_sum (int lq, int rq) {
        if (lb >= lq && rb <= rq)
            // если мы лежим полностью в отрезке запроса, вывести сумму
            return sum;
        if (max(lb, lq) >= min(rb, rq))
            // если мы не пересекаемся с отрезком запроса, вывести ноль
            return 0;
        // иначе всё сложно -- запускаемся от детей и пусть они там сами решают
        return l->get_sum(lq, rq) + r->get_sum(lq, rq);
    }
};

Посчитать число беспорядков в перестановке из \(n\) элементов (беспорядок или инверсия — это пара чисел \(i < j\), для которых \(p_i > p_j\)).

Эта задача решается просто, если уметь писать сортировку слиянием вручную. Но мы пойдем по другому пути. Создадим ДО для суммы на \(n\) элементов, изначально заполненное нулями. Теперь будем проходить по этому массиву слева направо. Когда обрабатываем очередное число \(x\), будем делать две вещи:

• Запросим сумму от \(k\) до \(n\) в ДО.
• Добавим единичку в \(k\)-тую позицию в ДО.

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

Даны \(n\) точек на плоскости с целыми координатами от 1до \(n\). Требуется ответить на \(m\) запросов количества точек на прямоугольнике.

Ответим на все запросы в оффлайн, используя метод сканирующей прямой:

• Разобьем запросы суммы на прямоугольнике на два запроса суммы на префиксах — сумма на прямоугольнике \([x_1, x_2] \times [y_1, y_2]\) равна сумме на прямоугольнике \([0, x_2] \times [y_1, y_2]\) минус сумма на прямоугольнике \([0, x_1] \times [y_1, y_2]\).
• Отсортируем теперь все точки и префиксные запросы по их \(x\). При этом, если у точки и запроса одинаковый \(x\), то точка должна идти раньше.
• Пройдёмся по ним в таком порядке и будем решать задачу для одномерной суммы: у нас есть операция «сделать +1 в \(y_i\)» и «вывести сумму с \(y_1\) по \(y_2\)».

Отложенные операции

Пусть теперь наш запрос обновления — это присвоение значения \(x\) всем элементам некоторого отрезка \([l, r)\), а не только одному.

Мы не хотим спускаться до каждого элемента, где меняется сумма — их может быть очень много. Мы схитрим, и при запросе присваивания будем, по возможности, помечать некоторые вершины, что они и все их дети «покрашены» в какое-то число. Непосредственно спускаться до листьев мы не будем.

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

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

Понятно, что от использования таких «запаздывающих» обновлений асимптотика никак не ухудшается, и мы можем всё так же решить задачу за \(O(n \log n)\).

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

struct segtree {
    int lb, rb;
    int sum = 0, assign = -1;
    segtree *l = 0, *r = 0;
    segtree (int _lb, int _rb) {
        lb = _lb, rb = _rb;
        if (lb + 1 < rb) {
            int t = (lb + rb) / 2;
            l = new segtree(lb, t);
            r = new segtree(t, rb);
        }
    }
    void push () {
        if (assign != -1) {
            sum = (rb-lb) * assign;
            if (l) { // если дети есть
                l->assign = assign;
                r->assign = assign;
            }
        }
        assign = -1;
    }
    void upd (int lq, int rq, int x) {
        push();
        if (lq <= lb && rb <= rq)
            assign = x;
        else if (l && max(lb, lq) < min(rb, rq)) {
            // если есть дети и отрезок запроса хоть как-то пересекается с нашим
            l->upd(lq, rq, x);
            r->upd(lq, rq, x);
            // ...дальше они сами разберутся
        }
    }
    int get_sum (int lq, int rq) {
        push();
        if (lb >= lq && rb <= rq)
            return sum;
        if (max(lb, lq) >= min(rb, rq))
            return 0;
        return l->get_sum(lq, rq) + r->get_sum(lq, rq);
    }
};

По-английски эта техника называется lazy propagation. Очень важно научиться её писать — она часто встречается на олимпиадах.

Идея «давайте будем всё делать в последний момент» применима не только в ДО, но и в других структурах и в реальной жизни.

Динамическое построение

А что, если у нас все индексы лежать не от в пределах \(10^5\), а, например, \(10^9\). Все асимптотики нас по прежнему устраивают (\(\log_2 10^6 \approx 20\), \(\log_2 10^9 \approx 30\)), кроме этапа построения.

Можно решить эту проблему так: откажемся от явного создания всех вершин дерева изначально. Изначально создадим только лишь корень, а остальные вершины будем создавать на ходу, когда в них потребуется записать что-то не дефолтное — как в lazy propagation.

Реализовать это можно так же, как и с push-ем: в начале всех методов будем проверять, что дети-вершины созданы, и создавать их, если это не так.

struct segtree {
    int lb, rb;
    int sum = 0;
    segtree *l = 0, *r = 0;
    segtree (int _lb, int _rb) {
        lb = _lb, rb = _rb;
        // а тут ничего нет
    }
    void extend () {
        if (!l && lb + 1 < rb) {
            int t = (lb + rb) / 2;
            l = new segtree(lb, t);
            r = new segtree(t, rb);
        }
    }
    void add (int k, int x) {
        extend();
        sum += x;
        if (l) {
            if (k < l->rb)
                l->add(k, x);
            else
                r->add(k, x);
        }
    }
    int get_sum (int lq, int rq) {
        if (lb >= lq && rb <= rq)
            return sum;
        if (max(lb, lq) >= min(rb, rq))
            return 0;
        extend();
        return l->get_sum(lq, rq) + r->get_sum(lq, rq);
    }
};

Но вообще, в большинстве случаев, использовать динамическое построение — это как стрелять из пушки по воробьям. Если все запросы известны заранее, то их координаты можно просто сжать перед обработкой запросов. Автор обычно делает это так:

vector<int> compress (vector<int> a) {
    vector<int> b = a;
    sort(b.begin(), b.end());
    b.erase(unique(b.begin(), b.end()), b.end());
    for (int &x : a) 
        x = int(lower_bound(b.begin(), b.end(), x) - b.begin());
    return a;
}

Персистентность

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

Известны персистентные версии многих структур: стека, очереди, СНМ, ДО. В случае со структурами данных на ссылках есть следующий общий подход: во всех методах, меняющих значения в вершинах, будем копировать ссылки на детей перед тем, как в них переходить и что-либо менять. Таким образом, мы всегда будем делать копию вершины перед тем, как что-либо менять в ней самой или её потомках. Вершины в момент \(t\) никогда не будут ссылаться на вершины, измененные после этого, и поэтому ничего не сломается.

У персистентных структур есть один минус: они обычно требуют больше памяти. В случае ДО мы будем создавать \(O(\log n)\) новых вершин на запрос, что означает общее потребление памяти \(O(m \log n)\).

struct segtree {
    int lb, rb;
    int sum = 0;
    segtree *l = 0, *r = 0;
    segtree (int _lb, int _rb) {
        lb = _lb, rb = _rb;
        if (lb != rb) {
            int t = (lb + rb) / 2;
            l = new segtree(lb, t);
            r = new segtree(t, rb);
        }
    }
    void copy () {
        if (l) {
            l = new segtree(*l);
            r = new segtree(*r);
        }
    }
    void add (int k, int x) {
        copy();
        sum += x;
        if (l) {
            if (k < l->rb) l->add(k, x);
            else r->add(k, x);
        }
    }
    int get_sum (int lq, int rq) {
        // этот метод ничего не меняет -- он и так хороший
        if (lq <= lb && rb <= rq)
            return sum;
        if (max(lb, lq) >= min(rb, rq))
            return 0;
        return l->get_sum(lq, rq) + r->get_sum(lq, rq);
    }
};

Даны \(n\) точек на плоскости. Нужно в онлайн ответить на \(q\) запросов суммы на прямоугольнике.

Если бы можно было отвечать в оффлайн, мы бы воспользовались методом сканирующей прямой — но так делать мы не можем. Вместо этого мы будем таким же образом добавлять точки в порядке увеличения \(x_i\) и декомпозировать запрос суммы на два, но при ответе на эти запросы мы будем доставать соответствующую версию ДО, которую мы получили, обработав нужное количество точек. Таким образом, можно отвечать на запросы в онлайн, но с \(O(n \log n)\) памяти.

Дан отрезок из \(n\) чисел от 1 до \(n\). Требуется ответить на \(q\) запросов \(k\)-той порядковой статистики на подотрезке.

Сделаем такой стандартный препроцессинг: пройдёмся с персистентным деревом отрезков для суммы по массиву. Когда будем обрабатывать элемент \(k\), добавим единицу к \(k\)-ому элементу.

Дальше определим разность деревьев как дерево отрезков, которое соответствует разности массивов. Заметим, что он неотрицательный. Его можно получить неявно, спускаясь одновременно в двух ДО и вместо sum использовать везде sum_r - sum_l.

Что будет находиться в разности \(r\)-го и \(l\)-го дерева? Там будут количества вхождений чисел на этом отрезке. В таком ДО не составить труда сделать спуск, который находит последнюю позицию, у которой сумма на соответствующем префиксе не превышает \(k\) — она и будет ответом.

Дан массив из \(n\) элементов. Требуется ответить на \(m\) запросов, есть ли на отрезке \([l, r]\) доминирующий элемент — тот, который встречается на нём хотя бы \(\frac{r-l}{2}\) раз.

У этой задачи есть удивительно простое решение — взять около 100 случайных элементов и каждый проверить, является ли он доминирующим (это можно проверить за \(O(\log n)\), посчитав для каждого значения отсортированный список позиций, на которых он встречается, и сделав два бинпоиска). Вероятность ошибки в худшем случае равна \(\frac{1}{2^{100}}\), и ей на практике можно пренебречь.

Но проверять 100 сэмплов — долго. Можно построить такое же ДО, как в прошлой задаче, и решать задачу «найти число, большее \(\frac{n}{2}\) в массиве на \(n\) элементов». Это тоже будет спуском по ДО: каждый раз идём в того сына, где сумма больше. Если в листе, куда мы пришли, значение больше нужного, возвращаем true, иначе false.

Другие реализации

Реализация на указателях проста и легко расширяема, но очень медленная и неэффективная по памяти: нужно хранить сами указатели, структура перестаёт помещаться в кэш, нужно много лишних раз ходить по ссылкам в память, только чтобы получить нужные вершины.

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

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

«ДО снизу». Можно делать все операции итеративно — так получится раз в 7 быстрее, но писать что-либо нетривиальное (например, массовые операции) так будет намного труднее. Подробнее смотрите в посте с CodeForces.