Treap

treap.cpp

#include <bits/stdc++.h>

// Легально скопипащено с https://github.com/kodo-pp/vsosh-region-preparation/

using namespace std;

default_random_engine gen;
uniform_real_distribution<double> dist(0.0, 1.0);

void init_rng()
{
	gen.seed(random_device()());
}

// Генерирует случайное число в промежутке [0.0, 1.0)
double random_double()
{
	return dist(gen);
}


// Функции и структуры для декартова дерева
namespace treap
{

// Вершина в декартовом дереве
struct Node
{
	// Конструктор: создаёт вершину с заданным ключом и случайным приоритетом
	Node(int x): x(x)
	{
		y = random_double();
	}

	Node* left = nullptr;   // Левый ребёнок
	Node* right = nullptr;  // Правый ребёнок
	int x;				  // Ключ
	int size = 1;
	double y;			   // Приоритет
};


int get_size(Node* a)
{
	if (a == nullptr) {
		return 0;
	} else {
		return a->size;
	}
}

void update(Node* a)
{
	if (a == nullptr) {
		return;
	}

	a->size = get_size(a->left) + get_size(a->right) + 1;
}

// Слияние двух деревьев
// Требования: дерево b полностью лежит строго правее, чем дерево a
Node* merge(Node* a, Node* b)
{
	// База рекурсии: merge(a, пустое дерево) = a
	// Пустое дерево обозначается нулевым указателем (nullptr)
	if (a == nullptr) {
		return b;
	}
	if (b == nullptr) {
		return a;
	}

	// Сравниваем приоритеты в корнях деревьев a и b
	if (a->y > b->y) {
		// Случай 1: в корне a приоритет больше, чем в корне b.
		// Значит корень дерева a будет корнем объединения деревьев a и b.
		// Левым потомком результирующего корня останется левое поддерево a (как и было до этого).
		// Правым потомком результирующего корня станет объединение правого поддерева a и дерева b.
		Node* m = merge(a->right, b);
		a->right = m;
		update(a);
		return a;
	} else {
		// Случай 2: в корне b приоритет больше, чем в корне a.
		// Случай противположен первому.
		// Корень дерева b будет корнем объединения деревьев a и b.
		// Правым потомком результирующего корня останется правое поддерево b.
		// Левым потомком результирующего корня станет объединение дерева a и левого поддерева a.
		Node* m = merge(a, b->left);
		b->left = m;
		update(b);
		return b;
	}
}

// Разделение дерева на два поддерева.
// Первое возвращаемое дерево будет иметь все ключи, меньшие или равные x,
// а второе — строго большие x.
pair<Node*, Node*> split(Node* a, int x)
{
	// База рекурсии: split(пустое дерево, x) = два пустых дерева (для любого x)
	if (a == nullptr) {
		return {nullptr, nullptr};
	}

	// Смотрим, в каком дереве будет лежать исходный корень
	if (a->x <= x) {
		// Первое (левое) поддерево
		// Значит, x лежит где-то справа от корня, и мы рекурсивно запускаемся от правого поддерева a
		auto [left_split, right_split] = split(a->right, x);
		a->right = left_split;
		update(a);
		return {a, right_split};
	} else {
		// Второе (правое) поддерево
		// Мне лень писать комментарии, помогите
		auto [left_split, right_split] = split(a->left, x);
		a->left = right_split;
		update(a);
		return {left_split, a};
	}
}

Node* add(Node* root, Node* v)
{
	if (root == nullptr) {
		return v;
	}
	if (v == nullptr) {
		throw std::runtime_error("Что вы тут нуллптрами кидаетесь?!");
	}
	if (root->y > v->y) {
		if (v->x <= root->x) {
			root->left = add(root->left, v);
		} else {
			root->right = add(root->right, v);
		}
		update(root);
		return root;
	} else {
		auto [l, r] = split(root, v->x);
		v->left = l;
		v->right = r;
		update(v);
		return v;
	}
}

// Добавляет элемент с заданным ключом к дереву
// Если передан nullptr, работает корректно
Node* add(Node* a, int x)
{
	return add(a, new Node(x));
}

// Выполняет поиск по ключу и возвращает true, если такой ключ существует, и false в противном случае
bool has_key(Node* a, int x)
{
	// Реально лень писать комментарии, но, короче, это просто поиск в бинарном дереве поиска
	if (a == nullptr) {
		return false;
	}
	if (a->x == x) {
		return true;
	} else if (a->x < x) {
		return has_key(a->right, x);
	} else /* a->x > x */ {
		return has_key(a->left, x);
	}
}


// Удаляет (под)дерево, освобождая память
void delete_tree(Node* a)
{
	if (a == nullptr) {
		return;
	}
	delete_tree(a->left);
	delete_tree(a->right);
	delete a;
}

// Удаляет все элементы с заданным ключом из дерева
Node* remove(Node* a, int x)
{
	// TODO: тру-путь

	// left_and_mid содержит элементы с ключом ≤ x
	// right — с ключом > x
	auto [left_and_mid, right] = split(a, x);

	// left содержит элементы с ключом < x
	// mid — с ключом == x
	auto [left, mid] = split(left_and_mid, x - 1); // TODO: работает только для целого ключа

	delete_tree(mid);

	auto full_tree = merge(left, right);
	return full_tree;
}

Node* kstat(Node* root, int k)
{
	if (root == nullptr) {
		return nullptr;
	}
	
	auto sz = get_size(root->left);
	if (sz == k) {
		return root;
	} else if (sz > k) {
		return kstat(root->left, k);
	} else {
		return kstat(root->right, k - sz - 1);
	}
}

} // namespace treap


void print(treap::Node* a, int nest = 0)
{
	string indent(nest * 4, ' ');
	if (a == nullptr) {
		cout << indent << "<empty tree>" << endl;
		return;
	}
	cout << indent << "x = " << a->x << ", " << "y = " << a->y << ", " << "size = " << a->size << endl;
	print(a->left, nest + 1);
	print(a->right, nest + 1);
}


int main()
{
	cout << boolalpha; // Вывод true/false вместо 1/0
	init_rng();

	treap::Node* root = nullptr;

	// Добавляем в ДД элементы (они могут повторяться)
	root = treap::add(root, 1);
	root = treap::add(root, 4);
	root = treap::add(root, 6);
	root = treap::add(root, 3);
	root = treap::add(root, 2);
	root = treap::add(root, 5);
	root = treap::add(root, 2);

	cout << "====== TREE AFTER INSERTION ======" << endl;
	print(root);
	cout << endl;

	// Несколько тестов
	cout << treap::has_key(root, 3) << endl;   // true
	cout << treap::has_key(root, 2) << endl;   // true
	cout << treap::has_key(root, -10) << endl; // false
	cout << treap::has_key(root, 8) << endl;   // false
	if (auto v = treap::kstat(root, 0); v == nullptr) {
		cout << "<not found>" << endl;
	} else {
		cout << v->x << endl;
	}
	if (auto v = treap::kstat(root, 1); v == nullptr) {
		cout << "<not found>" << endl;
	} else {
		cout << v->x << endl;
	}
	if (auto v = treap::kstat(root, 2); v == nullptr) {
		cout << "<not found>" << endl;
	} else {
		cout << v->x << endl;
	}
	if (auto v = treap::kstat(root, 40); v == nullptr) {
		cout << "<not found>" << endl;
	} else {
		cout << v->x << endl;
	}

	// Удаляем из ДД некоторые элементы
	root = treap::remove(root, 1);
	root = treap::remove(root, 2);
	cout << "====== TREE AFTER DELETION ======" << endl;
	print(root);
	cout << endl;

	// Несколько тестов
	cout << treap::has_key(root, 3) << endl;   // true
	cout << treap::has_key(root, 2) << endl;   // false
	cout << treap::has_key(root, -10) << endl; // false
	cout << treap::has_key(root, 8) << endl;   // false

	// Освобождаем память, выделенную под дерево
	treap::delete_tree(root);
}