Бинарные деревья поиска

Напомним: общая идея всех бинарных деревьев поиска - в том, что для заданного набора чисел $x_i$ строится дерево с этими числами в вершинах, обладающее следующими свойствами:

1. Оба поддерева корня являются бинарными деревьями поиска.
2. Все числа в левом поддереве не больше числа в корне.
3. Все числа в правом поддереве больше числа в корне.

Наличие такого дерева позволяет искать элемент в структуре, добавлять новые элементы или убирать имеющиеся (по сути, реализуется основной интерфейс std::set<int>). Главная проблема в том, что у такого дерева асимптотика на каждую операцию в худшем случае, например, когда в дерево последовательно добавляются вершины от 1 до n, составит $O(n)$. Для решения этой проблемы часто прибегают к разным модификациям, одну из которых мы и рассмотрим сегодня.

Декартово дерево

В декартовом дереве у каждой вершины есть не только ключ $x$, но и приоритет $y$ (из-за наличия двух координат, которые можно красиво нарисовать на декартовой плоскости, дерево и получило свое название). По ключу $x$ дерево является деревом поиска, а по $y$ - является кучей, то есть приоритет корня максимален.

Оказывается, если брать числа $x$ из входного набора данных, а $y$ генерировать случайным образом - дерево с помощью магии и теории вероятностей будет иметь среднюю глубину $O(\log n)$ и таким образом все операции будут выполняться за логарифм. Подробное описание данной магии выходит за рамки данного конспекта, поэтому читайте подробнее, например, здесь.

Пару слов про рандом

Здесь нам понадобится генерировать случайные числа и чем лучше они генерируются - тем более сбалансированным будет дерево. Стандартная функция для генерации целых чисел в C++ - rand(), но у нее есть серьезный минус: числа, которые она генерирует, малы, из-за этого весьма вероятны коллизии, что может влиять на качество дерева.

Поэтому часто используются альтернативные генераторы. Один из примеров - mt19937, используемый, например, следующим образом:

random_device rd;
mt19937_64 gen(rd());

Теперь функция gen() будет возвращать хорошо распределенные числа из всего диапозона 64-битных целых чисел. Если же достаточно 32 бит - можно убрать суффикс _64.

Реализация

Самый простой и понятный способ реализовать декартово дерево - использовать указатели.

Для вершин часто пишут структуру, имеющую примерно следующий вид:

struct Node {
    int x, y;
    Node* l = nullptr;
    Node* r = nullptr;
    Node (int val) {
        x = val;
        y = gen();
    }
};

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

Merge

Первой функцией будет merge - функция, принимающая два дерева, в которых в левом все ключи меньше, чем в правом.

Реализована она будет рекурсивно: сначала выберем корень нового дерева - это либо корень $tl$, либо корень $tr$. Нам нужно выбрать тот, у которого приоритет больше. Допустим, $tl->y > tr->y$, тогда новый корень - $tl$. Левым сыном нового корня будет прежний левый сын $tl$, а вот правым сыном будет $merge(tl->r, tr)$. Итого выходит примерно такой код:

Node* merge(Node* tl, Node* tr) {
    if (tl == nullptr) return tr;
    else if (tr == nullptr) return tl;
    if (tl->y > tr->y) {
        tl->r = merge(tl->r, tr);
        return tl;
    } else {
        tr->l = merge(tl, tr->l);
        return tr;
    }
}

Split

Эта функция будет принимать одно дерево $t$ и число $x$ и разделать дерево на два: в одном все значения не больше $x$, в другом - больше.

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

pair<node*, node*> split(node* t, int x) {
    if (t == nullptr) {
        return {nullptr, nullptr};
    }
    if (t->x <= x) {
        auto [tl, tr] = split(t->r, x);
        t->r = tl;
        return {t, tr};
    } else {
        auto [tl, tr] = split(t->l, x);
        t->l = tr;
        return {tl, t};
    }
}

Эти две функции сами по себе имеют малую ценность, но с помощью их комбинирования можно получать, что угодно. Например:

Insert

Эта функция будет добавлять ключ $x$ в дерево $t$ и возвращать новый указатель на корень дерева.

Node* insert(Node* t, int x) {
    auto [tl, tr] = split(t, x);
    Node *v = new Node(x);
    return merge(tl, merge(t, tr));
}

Сумма на отрезке

Внесём в Node ещё одно поле для суммы дерева - int sum. В split и merge будем поддерживать эту сумму, создав для этого вспомогательные функции:

void sum(Node* t) {
    if (!t) return 0;
    return t->sum;
}

void upd(Node* t) {
    if (!t) return;
    t->sum = sum(t->l) + sum(t->r) + t->x;
}

В split и merge необходимо добавить вызов upd от каждой возвращаемой вершины, например:

Node* merge(Node* tl, Node* tr) {
    if (tl == nullptr) return tr;
    else if (tr == nullptr) return tl;
    if (tl->y > tr->y) {
        tl->r = merge(tl->r, tr);
        upd(tl);
        return tl;
    } else {
        tr->l = merge(tl, tr->l);
        upd(tr);
        return tr;
    }
}

Теперь мы можем, например, посчитать сумму всех чисел между $l$ и $r$. Для этого - вырежем этот отрезок и вернём уже посчитанную сумму на нём:

pair<Node*, int> sum(Node* t, int l, int r) {
    auto [tmp, tr] = split(t, r);
    auto [tl, tmid] = split(tmp, l);
    int res = sum(tmid);
    return {merge(tl, merge(tmid, tr)), res};
}