洛谷P3369题解：Treap数据结构从入门到精通

19339201706

1. #include <iostream>
   
2. #include <cstdlib>
   
3. #include <climits>
   
4. using namespace std;
   
5. 
   
6. const int INF = 1e8; // 处理大数值范围
   
7. 
   
8. struct Node {
   
9.     int val, size, cnt;
   
10.     int priority;
    
11.     Node *l, *r;
    
12.     Node(int v) : val(v), size(1), cnt(1), l(nullptr), r(nullptr) {
    
13.         priority = rand() % INF; // 随机优先级保证平衡
    
14.     }
    
15. };
    
16. 
    
17. class Treap {
    
18. private:
    
19.     Node *root;
    
20.    
    
21.     // 更新节点子树大小
    
22.     void update(Node *node) {
    
23.         if(!node) return;
    
24.         node->size = node->cnt;
    
25.         if(node->l) node->size += node->l->size;
    
26.         if(node->r) node->size += node->r->size;
    
27.     }
    
28.    
    
29.     // 左旋操作
    
30.     void rotateLeft(Node *&node) {
    
31.         Node *temp = node->r;
    
32.         node->r = temp->l;
    
33.         temp->l = node;
    
34.         update(node);
    
35.         update(temp);
    
36.         node = temp;
    
37.     }
    
38.    
    
39.     // 右旋操作
    
40.     void rotateRight(Node *&node) {
    
41.         Node *temp = node->l;
    
42.         node->l = temp->r;
    
43.         temp->r = node;
    
44.         update(node);
    
45.         update(temp);
    
46.         node = temp;
    
47.     }
    
48.    
    
49.     // 插入操作
    
50.     void insert(Node *&node, int val) {
    
51.         if(!node) {
    
52.             node = new Node(val);
    
53.             return;
    
54.         }
    
55.         if(val == node->val) {
    
56.             node->cnt++; // 重复值计数
    
57.         } else if(val < node->val) {
    
58.             insert(node->l, val);
    
59.             if(node->l->priority > node->priority)
    
60.                 rotateRight(node); // 维护堆性质
    
61.         } else {
    
62.             insert(node->r, val);
    
63.             if(node->r->priority > node->priority)
    
64.                 rotateLeft(node); // 维护堆性质
    
65.         }
    
66.         update(node);
    
67.     }
    
68.    
    
69.     // 删除操作
    
70.     void remove(Node *&node, int val) {
    
71.         if(!node) return;
    
72.         if(val < node->val) {
    
73.             remove(node->l, val);
    
74.         } else if(val > node->val) {
    
75.             remove(node->r, val);
    
76.         } else {
    
77.             if(node->cnt > 1) {
    
78.                 node->cnt--; // 减少计数
    
79.             } else {
    
80.                 if(!node->l || !node->r) {
    
81.                     Node *temp = node->l ? node->l : node->r;
    
82.                     delete node;
    
83.                     node = temp; // 单子树情况直接替换
    
84.                 } else {
    
85.                     // 选择优先级高的子树旋转
    
86.                     if(node->l->priority > node->r->priority) {
    
87.                         rotateRight(node);
    
88.                         remove(node->r, val);
    
89.                     } else {
    
90.                         rotateLeft(node);
    
91.                         remove(node->l, val);
    
92.                     }
    
93.                 }
    
94.             }
    
95.         }
    
96.         if(node) update(node);
    
97.     }
    
98.    
    
99.     // 获取排名
    
100.     int getRank(Node *node, int val) {
     
101.         if(!node) return 0;
     
102.         if(val < node->val) return getRank(node->l, val);
     
103.         int leftSize = node->l ? node->l->size : 0;
     
104.         if(val == node->val) return leftSize + 1;
     
105.         return leftSize + node->cnt + getRank(node->r, val);
     
106.     }
     
107.    
     
108.     // 根据排名获取值
     
109.     int getValue(Node *node, int rank) {
     
110.         if(!node) return INF;
     
111.         int leftSize = node->l ? node->l->size : 0;
     
112.         if(rank <= leftSize) return getValue(node->l, rank);
     
113.         if(rank <= leftSize + node->cnt) return node->val;
     
114.         return getValue(node->r, rank - leftSize - node->cnt);
     
115.     }
     
116.    
     
117.     // 获取前驱
     
118.     int getPre(Node *node, int val) {
     
119.         if(!node) return -INF;
     
120.         if(node->val >= val) return getPre(node->l, val);
     
121.         return max(node->val, getPre(node->r, val));
     
122.     }
     
123.    
     
124.     // 获取后继
     
125.     int getNext(Node *node, int val) {
     
126.         if(!node) return INF;
     
127.         if(node->val <= val) return getNext(node->r, val);
     
128.         return min(node->val, getNext(node->l, val));
     
129.     }
     
130. 
     
131. public:
     
132.     Treap() : root(nullptr) { srand(time(0)); }
     
133.    
     
134.     // 公开接口
     
135.     void insert(int val) { insert(root, val); }
     
136.     void remove(int val) { remove(root, val); }
     
137.     int getRank(int val) { return getRank(root, val); }
     
138.     int getValue(int rank) { return getValue(root, rank); }
     
139.     int getPre(int val) { return getPre(root, val); }
     
140.     int getNext(int val) { return getNext(root, val); }
     
141. };
     
142. 
     
143. int main() {
     
144.     ios::sync_with_stdio(false);
     
145.     cin.tie(0);
     
146.    
     
147.     Treap treap;
     
148.     int n, opt, x;
     
149.     cin >> n;
     
150.     while(n--) {
     
151.         cin >> opt >> x;
     
152.         switch(opt) {
     
153.             case 1: treap.insert(x); break;
     
154.             case 2: treap.remove(x); break;
     
155.             case 3: cout << treap.getRank(x) << '\n'; break;
     
156.             case 4: cout << treap.getValue(x) << '\n'; break;
     
157.             case 5: cout << treap.getPre(x) << '\n'; break;
     
158.             case 6: cout << treap.getNext(x) << '\n'; break;
     
159.         }
     
160.     }
     
161.     return 0;
     
162. }

复制代码

来源：洛谷题解