Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <iostream>
- #include <vector>
- #include <fstream>
- #include <algorithm>
- #include <set>
- #include <chrono>
- #include <random>
- #define len(v) (int)v.size()
- #define all(v) v.begin(), v.end()
- #define rall(v) v.rbegin(), v.rend()
- using namespace std;
- const int maxn = 1e5 + 10;
- const int inf = 1e9;
- const int m = 1e9 + 7;
- mt19937 rnd(chrono::high_resolution_clock::now().time_since_epoch().count());
- struct Node
- {
- int x; // Хранимый ключ
- int y; // "Порядок" вершины
- int sz; // Размер поддерева (количество узлов, находящихся ниже данного)
- Node* l, * r; // Указатели на потомков данного узла
- // Конструктор узла (объявлен в описании самой структуры для простоты)
- Node(int _x)
- {
- x = _x;
- y = rnd() % m;
- 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);
- }
- }
- bool find(Node* t, int k) {
- if (!t)
- return false;
- if (t->x == k)
- return true;
- if (t->x > k)
- return find(t->l, k);
- else
- return find(t->r, k);
- }
- Node* next(Node* t, int k) {
- Node *cur = t, *succ = NULL;
- while (cur != NULL) {
- if (cur->x > k) {
- succ = cur;
- cur = cur->l;
- }
- else
- cur = cur->r;
- }
- return succ;
- }
- Node* prev(Node* t, int k, int res = -1) {
- Node* cur = t, * succ = NULL;
- while (cur != NULL) {
- if (cur->x < k) {
- succ = cur;
- cur = cur->r;
- }
- else
- cur = cur->l;
- }
- return succ;
- }
- void solve() {
- Node* treap = nullptr;
- string cmd;
- while (cin >> cmd) {
- int x; cin >> x;
- if (cmd[0] == 'i') {
- add(treap, x);
- }
- else if (cmd[0] == 'd') {
- remove(treap, x);
- }
- else if (cmd[0] == 'e') {
- bool res = find(treap, x);
- if (res)
- cout << "true\n";
- else
- cout << "false\n";
- }
- else if (cmd[0] == 'n') {
- Node* res = next(treap, x);
- if (res == NULL)
- cout << "none\n";
- else
- cout << res->x << "\n";
- }
- else if (cmd[0] == 'p') {
- Node* res = prev(treap, x);
- if (res == NULL)
- cout << "none\n";
- else
- cout << res->x << "\n";
- }
- }
- }
- int main() {
- ios::sync_with_stdio(false);
- cin.tie(nullptr);
- cout.tie(nullptr);
- solve();
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement