Декартово дерево по явному ключу

treap_explicit.cpp

#include <iostream>
#include <cstdlib>
#include <vector>

// --- Декартово дерево (примерно так, как это было на занятии)

using namespace std;

// Определение узла декартового дерева
struct Node
{
    int x; // Хранимый ключ
    int y; // "Порядок" вершины
    int sz; // Размер поддерева (количество узлов, находящихся ниже данного)
    Node *l, *r; // Указатели на потомков данного узла

    // Конструктор узла (объявлен в описании самой структуры для простоты)
    Node(int _x)
    {
        x = _x;
        y = rand();
        sz = 1;
        l = r = nullptr;
    }
};

// Прототипы функций для работы с размером поддерева - их реализуем ниже.
int get_sz(Node *t);
void update(Node *t);

// --- Основные операции - слияние (merge) и разделение (split) деревьев
// Важно! При выполнении этих операций исходные деревья могут перестать быть
// валидными!

// Слияние двух деревьев (t1, t2). Работает корректно тогда и только тогда,
// когда max(t1) < min(t2). Результат - корень нового дерева.
Node* merge(Node *t1, Node *t2)
{
    // Пустое дерево - то, у которого корень - nullptr.
    // Считаем, что слияние дерева с пустым деревом - это исходное дерево.
    if (t1 == nullptr) { return t2; }
    if (t2 == nullptr) { return t1; }

    // Дерево с большим значением y становится новым корнем
    if (t1->y > t2->y)
    {
        // Сливаем правое поддерево первого дерева со вторым деревом и ставим
        // результат в правое поддерево первого дерева.
        t1->r = merge(t1->r, t2);
        // Поскольку теперь дерево t1 поменялось, запускаем в нём обновление.
        update(t1);
        // Новый корень дерева - t1.
        return t1;
    }
    else
    {
        // Сливаем всё первое дерево с левым поддеревом второго дерева
        // (обязательно именно в таком порядке!!) и ставим результат в левое поддерево
        // второго дерева.
        t2->l = merge(t1, t2->l);
        // Пересчитываем значение в новом корне дерева
        update(t2);
        return t2;
    }
}

// Разделение дерева на два по заданному ключу. После этого в первом дереве
// будут значения (-inf, x), во втором - [x, +inf)
void split(Node *t, int x, Node *&t1, Node *&t2)
{
    // Пустое дерево будет в любом случае разделено на два пустых поддерева
    if (t == nullptr)
    {
        t1 = t2 = nullptr;
        return;
    }
    // Если текущая вершина содержит значение, меньшее, чем заданное,
    // то она отправляется в первое дерево. Правое поддерево данной вершины
    // в принципе может оказаться больше или равно x, так что его тоже разрезаем,
    // и записываем первое дерево-результат на его место.
    if (t->x < x)
    {
        split(t->r, x, t->r, t2);
        t1 = t;
    }
    else
    {
        // В противном случае "режем" левое поддерево. Рассуждения остаются такими
        // же (текущая вершина отправляется во второе дерево-результат)
        split(t->l, x, t1, t->l);
        t2 = t;
    }
    // В процессе отрезания мы модифицируем дерево, поэтому для данной вершины надо
    // пересчитать хранимое выражение
    update(t);
}

// Безопасное получение размера поддерева
int get_sz(Node *t)
{
    // Размер пустых деревьев считаем равным нулю
    if (t == nullptr) { return 0; }
    // У непустых поддеревьев размер хранится в корне
    return t->sz;
}

// Обновление размера поддерева
void update(Node *t)
{
    // Размер обновляем только для непустых деревьев
    if (t != nullptr)
    {
        t->sz = 1 + get_sz(t->l) + get_sz(t->r);
    }
}


// --- Операции, проводимые с деревом

// Добавление нового элемента x в дерево t. Корень дерева может измениться!
void add(Node *&t, int x)
{
    Node *t1, *t2;
    // Для добавления делаем следующее:
    // - Разрезаем исходное дерево по ключу x. В левом поддереве все элементы меньше x,
    //   в правом - не меньше.
    split(t, x, t1, t2);
    // - Создаём новое дерево из одной вершины - собственно, x.
    Node* new_tree = new Node(x);
    // - Производим слияние левого поддерева с новым, потом слияние результата с правым
    //   Результат слияния - новый корень дерева.
    t = merge(merge(t1, new_tree), t2);
}

// Удаление вершины из дерева. В данном случае работать будет только с целыми числами!
void remove(Node *&t, int x)
{
    Node *t1, *t2, *t3, *t4;
    // Для удаления делаем следующее:
    // - Разрезаем исходное дерево по ключу x.
    split(t, x, t1, t2);
    // - Разрезаем правое поддерево по ключу x + 1 (вот зачем нужны были целые числа!)
    split(t2, x + 1, t3, t4);
    // - Соединяем деревья t1 и t4, которые теперь не содержат ключа x
    //   (он остался в дереве t3)
    t = merge(t1, t4);
    // - Очищаем память, занимаемую деревом t3 (если не хотим утечек)
    delete t3;
}

// Вывод элементов дерева на экран в отсортированном порядке (обход ЛКП)
void print(Node *t)
{
    // У пустых деревьев выводить нечего
    if (t != nullptr)
    {
        // Сначала выводим всё из левого поддерева, затем то, что хранится в
        // корне, затем всё из правого поддерева
        print(t->l);
        cout << t->x << " ";
        print(t->r);
    }
}

// Получение элемента, стоящего на k-м месте в дереве (начиная с единицы!)
int get_k(Node *t, int k)
{
    if (k < get_sz(t->l))
    {
        return get_k(t->l, k);
    }
    else if (k == get_sz(t->l))
    {
        return t->x;
    }
    else
    {
        return get_k(t->r, k - get_sz(t->l) - 1);
    }
}

// Демонстрация операций с деревом
int main() {
    Node *treap = nullptr;
    vector<int> v{5, 1, 3, 8, 9, 6, 2};
    for(const auto &e : v)
    {
        add(treap, e);
    }
    print(treap);
    cout << endl;
    for(int i = 0; i < v.size(); ++i)
    {
        cout << get_k(treap, i) << " ";
    }
    cout << endl;
    remove(treap, 5);
    print(treap);
    while(treap != nullptr)
    {
        remove(treap, get_k(treap, 0));
        print(treap);
        cout << endl;
    }
    return 0;
}