Объяснение

Декартово дерево, или Treap, — это структура данных, которая соединяет две идеи:

• Binary Search Tree — по значению key
• Heap — по случайному приоритету prior

Благодаря случайным приоритетам дерево обычно получается сбалансированным. Поэтому основные операции работают за O(log n) в среднем.

Что храним в массиве?

struct Node {
    int key, prior;
    int l, r;
    int cnt;
} tr[N];

int root;

• key — значение вершины
• prior — случайный приоритет
• l — индекс левого сына
• r — индекс правого сына
• cnt — размер поддерева
• root — корень всего дерева

Главные операции

split

split(v, x) делит дерево v на две части:

• первая часть — все значения <= x
• вторая часть — все значения > x

merge

merge(a, b) соединяет два дерева в одно.

Пример задачи для реализации

Нужно хранить множество чисел и обрабатывать команды.

Input

10
insert 10
insert 5
insert 20
find 10
find 7
kth 2
erase 10
find 10
kth 2
find 20

Output

YES
NO
10
NO
20
YES

Полная реализация

#include <bits/stdc++.h>
using namespace std;

const int N = 100010;

struct Node {
    int key, prior;
    int l, r;
    int cnt;
} tr[N];

int root;
int idx;

int sizeOf(int v) {
    if (v == 0) {
        return 0;
    }

    return tr[v].cnt;
}

void pull(int v) {
    if (v == 0) {
        return;
    }

    tr[v].cnt = 1 + sizeOf(tr[v].l) + sizeOf(tr[v].r);
}

int newNode(int x) {
    idx++;

    tr[idx].key = x;
    tr[idx].prior = rand();
    tr[idx].l = 0;
    tr[idx].r = 0;
    tr[idx].cnt = 1;

    return idx;
}

pair<int, int> split(int v, int x) {
    if (v == 0) {
        return {0, 0};
    }

    if (tr[v].key <= x) {
        pair<int, int> p = split(tr[v].r, x);

        tr[v].r = p.first;
        pull(v);

        return {v, p.second};
    } else {
        pair<int, int> p = split(tr[v].l, x);

        tr[v].l = p.second;
        pull(v);

        return {p.first, v};
    }
}

int merge(int a, int b) {
    if (a == 0) {
        return b;
    }

    if (b == 0) {
        return a;
    }

    if (tr[a].prior > tr[b].prior) {
        tr[a].r = merge(tr[a].r, b);
        pull(a);

        return a;
    } else {
        tr[b].l = merge(a, tr[b].l);
        pull(b);

        return b;
    }
}

bool findKey(int v, int x) {
    if (v == 0) {
        return false;
    }

    if (tr[v].key == x) {
        return true;
    }

    if (x < tr[v].key) {
        return findKey(tr[v].l, x);
    } else {
        return findKey(tr[v].r, x);
    }
}

void insertKey(int x) {
    if (findKey(root, x)) {
        return;
    }

    pair<int, int> p = split(root, x);

    int leftTree = p.first;
    int rightTree = p.second;
    int node = newNode(x);

    root = merge(merge(leftTree, node), rightTree);
}

void eraseKey(int &v, int x) {
    if (v == 0) {
        return;
    }

    if (tr[v].key == x) {
        v = merge(tr[v].l, tr[v].r);
        return;
    }

    if (x < tr[v].key) {
        eraseKey(tr[v].l, x);
    } else {
        eraseKey(tr[v].r, x);
    }

    pull(v);
}

int kth(int v, int k) {
    int leftSize = sizeOf(tr[v].l);

    if (k == leftSize + 1) {
        return tr[v].key;
    }

    if (k <= leftSize) {
        return kth(tr[v].l, k);
    }

    return kth(tr[v].r, k - leftSize - 1);
}

void printTree(int v) {
    if (v == 0) {
        return;
    }

    printTree(tr[v].l);
    cout << tr[v].key << ' ';
    printTree(tr[v].r);
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    srand(time(0));

    int q;
    cin >> q;

    while (q--) {
        string op;
        int x;

        cin >> op >> x;

        if (op == "insert") {
            insertKey(x);
        } else if (op == "erase") {
            eraseKey(root, x);
        } else if (op == "find") {
            if (findKey(root, x)) {
                cout << "YES\n";
            } else {
                cout << "NO\n";
            }
        } else if (op == "kth") {
            if (x >= 1 && x <= sizeOf(root)) {
                cout << kth(root, x) << '\n';
            }
        }
    }

    return 0;
}

Объяснение функций

1. sizeOf

int sizeOf(int v) {
    if (v == 0) {
        return 0;
    }

    return tr[v].cnt;
}

Функция возвращает размер поддерева. Если вершина пустая, размер равен 0.

2. pull

void pull(int v) {
    tr[v].cnt = 1 + sizeOf(tr[v].l) + sizeOf(tr[v].r);
}

Функция пересчитывает размер поддерева вершины v.

3. newNode

int newNode(int x)

Создает новую вершину в массиве tr. Возвращает индекс этой вершины.

4. split

pair<int, int> split(int v, int x)

Делит дерево на две части:

• первая часть содержит значения <= x
• вторая часть содержит значения > x

Если текущий key <= x, значит эта вершина должна попасть в левую часть. Тогда мы делим её правого сына.

if (tr[v].key <= x) {
    pair<int, int> p = split(tr[v].r, x);

    tr[v].r = p.first;
    pull(v);

    return {v, p.second};
}

Если текущий key > x, значит эта вершина должна попасть в правую часть. Тогда мы делим её левого сына.

else {
    pair<int, int> p = split(tr[v].l, x);

    tr[v].l = p.second;
    pull(v);

    return {p.first, v};
}

5. merge

int merge(int a, int b)

Соединяет два дерева. Все значения в a должны быть меньше всех значений в b.

Корнем становится дерево с большим prior.

if (tr[a].prior > tr[b].prior) {
    tr[a].r = merge(tr[a].r, b);
    pull(a);

    return a;
}

Если у b приоритет больше, то b становится корнем.

else {
    tr[b].l = merge(a, tr[b].l);
    pull(b);

    return b;
}

6. findKey

bool findKey(int v, int x)

Ищет число x как в обычном бинарном дереве поиска.

7. insertKey

void insertKey(int x)

Добавляет число x в дерево.

pair<int, int> p = split(root, x);

int leftTree = p.first;
int rightTree = p.second;
int node = newNode(x);

root = merge(merge(leftTree, node), rightTree);

8. eraseKey

void eraseKey(int &v, int x)

Удаляет число x из дерева. Когда нашли нужную вершину, заменяем её на merge левого и правого сына.

if (tr[v].key == x) {
    v = merge(tr[v].l, tr[v].r);
    return;
}

9. kth

int kth(int v, int k)

Находит k-й элемент по возрастанию. Для этого используется размер левого поддерева.

Если слева стоит leftSize элементов:

• если k == leftSize + 1, ответ — текущая вершина
• если k <= leftSize, идем влево
• иначе идем вправо и уменьшаем k

Асимптотика