Treap是一种非常好写而且易于理解的平衡树,然而效率基本看脸,毕竟用到了随机算法。
期望情况下 插入 删除等操作的复杂度均为$$O(log n)$$
详情见注释
#include
using std::cin;
using std::cout;
using std::endl;
using std::string;
using int_t = long long int;
//表示treap中的一个节点
struct Node {
//该节点的值
int_t value;
//该节点的随机优先级,用于将treap组织成一个小根堆
//但堆是完全平衡的 而二叉搜索树显然不容易完全平衡 所以priority仅用来使treap带有堆的性质
//所以treap仍然是二叉搜索树
//某个节点的priority必须必它的子节点的小
int_t priority;
//同样值的个数,用于值会重复的情况
int_t count = 1;
//以该节点为根节点的子树的节点数(考虑重复值的情况)
int_t size = 1;
//左子节点
Node* left = nullptr;
//右子节点
Node* right = nullptr;
//父节点 注意 这个值只会在查询时更新
Node* parent = nullptr;
Node(int_t value) :
value(value) {
priority = rand();
}
//重新计算该节点的size值
void maintain() {
size = count;
if (left) size += left->size;
if (right) size += right->size;
}
//返回当前节点左子节点的size 即比当前节点小的节点的个数
int_t leftSize() {
if (left) return left->size;
else return 0;
}
};
Node* root = nullptr;
//在treap中搜索一个节点,返回节点的指针,如果没找到就返回nullptr
Node* lookup(Node* node, int_t value, Node* from = nullptr) {
if (node == nullptr) {
return nullptr;
}
//搜索的过程中顺便设置一下父节点
node->parent = from;
if (node->value == value) {
return node;
} else if (value < node->value) {
return lookup(node->left, value, node);
} else {
return lookup(node->right, value, node);
}
}
//把node的右子节点绕node左旋
void leftRotate(Node*& node) {
//把node的右节点存起来
Node* temp = node->right;
//然后把node的右节点换成tmp的左节点
node->right = temp->left;
//把temp的左节点换成node
temp->left = node;
//然后重新算一下size值
node->maintain();
temp->maintain();
node = temp;
}
//把node的左子节点绕node右旋
void rightRotate(Node* & node) {
Node* temp = node->left;
node->left = temp->right;
temp->right = node;
node->maintain();
temp->maintain();
node = temp;
}
//向以node为根节点的子树中插入value 并且维护最小堆
void insert(Node*& node, int_t value) {
if (node == nullptr) {
//如果node==nullptr就说明应该在这个位置插入节点
node = new Node(value);
} else {
//向左子树插入
if (value < node->value) {
insert(node->left, value);
//插入结束后,如果插入的节点破坏了最小堆的性质,就需要旋转一下,把优先级小的节点转到上面
//这样做不会破坏二叉搜索树的性质
if (node->left->priority > node->priority) {
rightRotate(node);
}
} else if (value > node->value) {
insert(node->right, value);
if (node->right->priority > node->priority) {
leftRotate(node);
}
} else {
//特殊情况 如果value已经存在于树中 则直接在计数器+1
node->count++;
}
}
//重新算一下size
if (node) node->maintain();
}
//从以node为根节点的子树中删除value
void remove(Node* & node, int_t value) {
//找到了要删除的节点
if (node->value == value) {
//特殊情况 如果有多个值相同的节点 那么计数器-1即可
if (node->count > 1) {
node->count--;
} else if (node->left == nullptr) {
//如果当前节点没有左子节点 那么就直接把右子节点接上来
Node* bak = node->right;
delete node;
node = bak;
} else if (node->right == nullptr) {
Node* bak = node->left;
delete node;
node = bak;
} else {
//如果node有两棵子树,那么就把优先级小的那一棵旋转到node,然后在另一个子树(此时node在这个子树里)删除node
//执行完多次旋转后,最后一定会到node只有一棵子树或者没有子树的情况
if (node->left->priority < node->right->priority) {
rightRotate(node);
remove(node->right, value);
} else {
leftRotate(node);
remove(node->left, value);
}
}
} else {
if (value < node->value) {
remove(node->left, value);
} else {
remove(node->right, value);
}
}
if (node) node->maintain();
}
//输出先序遍历
void preOrder(Node* node) {
if (node == nullptr) return;
cout << node->value << endl;
preOrder(node->left);
preOrder(node->right);
}
//输出中序遍历
void inOrder(Node* node) {
if (node == nullptr) return;
inOrder(node->left);
cout << node->value << endl;
inOrder(node->right);
}
//输出后序遍历
void postOrder(Node* node) {
if (node == nullptr) return;
postOrder(node->left);
postOrder(node->right);
cout << node->value << endl;
}
//求某个节点的前驱 即比这个节点小的最大的节点
//返回该节点指针 若不存在则返回nullptr
//需要保证value一定存在于树中
Node* pre(Node* node, int_t value) {
node = lookup(node, value);
if (node->left) {
//如果node有左子节点 那么node的前驱一定是node的左子节点的最右边的节点
node = node->left;
while (node->right) node = node->right;
return node;
} else {
//如果node没有左子节点 则顺着node一直往上找 直到找到node是某个节点的左子节点为止
Node* parent = node->parent;
while (parent && node != parent->right) {
node = parent;
parent = parent->parent;
}
return parent;
}
}
//求节点的后继 与前驱完全对称 不再赘述
Node* post(Node* node, int_t value) {
node = lookup(node, value);
if (node->right) {
node = node->right;
while (node->left) node = node->left;
return node;
} else {
Node* parent = node->parent;
while (node && node != parent->left) {
node = parent;
parent = parent->parent;
}
return parent;
}
}
//在以node为根节点的子树中查询value的排名
//要求value必须存在于树中
int_t rank(Node* node, int_t value) {
if (node == nullptr) return 0;
//如果value恰好等于node->value 就说明node的左子树的所有节点恰好比value小 node的右子树的节点恰好比value 大 value的名次就是node左子树的节点数
if (value == node->value) {
return node->leftSize();
} else if (value < node->value) {
//如果value比node->value小 那么就应该去node的左子树确定value的名次
return rank(node->left, value);
} else {
//如果value比node大 那么value的名次就是 node左子树的大小+值为node->value的元素个数+value在右子树的排名
return rank(node->right, value) + node->count + node->leftSize();
}
}
//在以node为根的子树内找第k小的元素
Node* kth(Node* node, int_t k) {
if (node == nullptr || k < 1 || k > node->size) return nullptr;
//如果k恰好是node的名次
//因为同一个元素可能有多个 所以需要这样写
if (node->leftSize() + 1 <= k && k <= node->leftSize() + node->count) {
return node;
} else if (k <= node->leftSize()) {
return kth(node->left, k);
} else {
return kth(node->right, k - node->leftSize() - node->count);
}
}
int main() {
// freopen("Z:/data.in", "r", stdin);
// freopen("Z:/data.out", "w", stdout);
int_t n;
cin>>n;
while (n--) {
string opt;
cin>>opt;
if (opt == "insert" || opt == "1") {
int_t x;
cin>>x;
insert(root, x);
} else if (opt == "find") {
int_t x;
cin>>x;
cout << lookup(root, x) << endl;
} else if (opt == "remove" || opt == "2") {
int_t x;
cin>>x;
remove(root, x);
} else if (opt == "pre") {
preOrder(root);
} else if (opt == "in") {
inOrder(root);
} else if (opt == "post") {
postOrder(root);
} else if (opt == "prev" || opt == "5") {
int_t x;
cin>>x;
insert(root, x);
cout << pre(root, x)->value << endl;
remove(root, x);
} else if (opt == "next" || opt == "6") {
int_t x;
cin>>x;
insert(root, x);
cout << post(root, x)->value << endl;
remove(root, x);
} else if (opt == "rank" || opt == "3") {
int_t x;
cin>>x;
insert(root, x);
cout << rank(root, x) + 1 << endl;
remove(root, x);
} else if (opt == "kth" || opt == "4") {
int_t x;
cin>>x;
Node* node = kth(root, x);
if (node)
cout << node->value << endl;
else cout << 0 << endl;
}
}
}
