洛谷P3369题解：Treap数据结构从入门到精通
https://dajuwangluo.cn/zb_users/upload/2025/07/202507111752229061585630.jpg#include <iostream>
#include <cstdlib>
#include <climits>
using namespace std;
const int INF = 1e8; // 处理大数值范围
struct Node {
int val, size, cnt;
int priority;
Node *l, *r;
Node(int v) : val(v), size(1), cnt(1), l(nullptr), r(nullptr) {
priority = rand() % INF; // 随机优先级保证平衡
}
};
class Treap {
private:
Node *root;
// 更新节点子树大小
void update(Node *node) {
if(!node) return;
node->size = node->cnt;
if(node->l) node->size += node->l->size;
if(node->r) node->size += node->r->size;
}
// 左旋操作
void rotateLeft(Node *&node) {
Node *temp = node->r;
node->r = temp->l;
temp->l = node;
update(node);
update(temp);
node = temp;
}
// 右旋操作
void rotateRight(Node *&node) {
Node *temp = node->l;
node->l = temp->r;
temp->r = node;
update(node);
update(temp);
node = temp;
}
// 插入操作
void insert(Node *&node, int val) {
if(!node) {
node = new Node(val);
return;
}
if(val == node->val) {
node->cnt++; // 重复值计数
} else if(val < node->val) {
insert(node->l, val);
if(node->l->priority > node->priority)
rotateRight(node); // 维护堆性质
} else {
insert(node->r, val);
if(node->r->priority > node->priority)
rotateLeft(node); // 维护堆性质
}
update(node);
}
// 删除操作
void remove(Node *&node, int val) {
if(!node) return;
if(val < node->val) {
remove(node->l, val);
} else if(val > node->val) {
remove(node->r, val);
} else {
if(node->cnt > 1) {
node->cnt--; // 减少计数
} else {
if(!node->l || !node->r) {
Node *temp = node->l ? node->l : node->r;
delete node;
node = temp; // 单子树情况直接替换
} else {
// 选择优先级高的子树旋转
if(node->l->priority > node->r->priority) {
rotateRight(node);
remove(node->r, val);
} else {
rotateLeft(node);
remove(node->l, val);
}
}
}
}
if(node) update(node);
}
// 获取排名
int getRank(Node *node, int val) {
if(!node) return 0;
if(val < node->val) return getRank(node->l, val);
int leftSize = node->l ? node->l->size : 0;
if(val == node->val) return leftSize + 1;
return leftSize + node->cnt + getRank(node->r, val);
}
// 根据排名获取值
int getValue(Node *node, int rank) {
if(!node) return INF;
int leftSize = node->l ? node->l->size : 0;
if(rank <= leftSize) return getValue(node->l, rank);
if(rank <= leftSize + node->cnt) return node->val;
return getValue(node->r, rank - leftSize - node->cnt);
}
// 获取前驱
int getPre(Node *node, int val) {
if(!node) return -INF;
if(node->val >= val) return getPre(node->l, val);
return max(node->val, getPre(node->r, val));
}
// 获取后继
int getNext(Node *node, int val) {
if(!node) return INF;
if(node->val <= val) return getNext(node->r, val);
return min(node->val, getNext(node->l, val));
}
public:
Treap() : root(nullptr) { srand(time(0)); }
// 公开接口
void insert(int val) { insert(root, val); }
void remove(int val) { remove(root, val); }
int getRank(int val) { return getRank(root, val); }
int getValue(int rank) { return getValue(root, rank); }
int getPre(int val) { return getPre(root, val); }
int getNext(int val) { return getNext(root, val); }
};
int main() {
ios::sync_with_stdio(false);
cin.tie(0);
Treap treap;
int n, opt, x;
cin >> n;
while(n--) {
cin >> opt >> x;
switch(opt) {
case 1: treap.insert(x); break;
case 2: treap.remove(x); break;
case 3: cout << treap.getRank(x) << '\n'; break;
case 4: cout << treap.getValue(x) << '\n'; break;
case 5: cout << treap.getPre(x) << '\n'; break;
case 6: cout << treap.getNext(x) << '\n'; break;
}
}
return 0;
}
来源：洛谷题解
页:
[1]