BST基础教学（详细注释）

  1 #include<iostream>
  2 #include<cstdio>
  3 #include<queue>
  4 #include<cstring>
  5 #include<algorithm>
  6 #define re return
  7 #define rep(i,a,n) for(int i = a;i <= n;i++)
  8 #define per(i,n,a) for(int i = n;i >= a;i--)
  9 typedef long long LL;
 10 using namespace std;
 11 int read() {
 12     int out = 0,flag = 1;
 13     char c = getchar();
 14     while(c < ‘0‘ || c > ‘9‘) {
 15         if(c == ‘-‘) flag = -1;
 16         c = getchar();
 17     }
 18     while(c >= ‘0‘ && c <= ‘9‘) {
 19         out = out * 10 + c - ‘0‘;
 20         c = getchar();
 21     }
 22     re flag * out;
 23 }
 24 //head
 25 
 26 const int maxn = 1000019,INF = 1e9;
 27 int na;
 28 int ch[maxn][2];//[i][0]代表i左儿子，[i][1]代表i右儿子
 29 int val[maxn],dat[maxn];
 30 int size[maxn],cnt[maxn];
 31 int tot,root;
 32 int New(int v) { //新增节点，
 33     val[++tot] = v;//节点赋值
 34     dat[tot] = rand();//随机优先级
 35     size[tot] = 1;//目前是新建叶子节点，所以子树大小为1
 36     cnt[tot] = 1;//新建节点同理副本数为1
 37     return tot;
 38 }
 39 void pushup(int id) { //和线段树的pushup更新一样
 40     size[id] = size[ch[id][0]] + size[ch[id][1]] + cnt[id];
 41     //本节点子树大小 = 左儿子子树大小 + 右儿子子树大小 + 本节点副本数
 42 }
 43 void build() {
 44     root = New(-INF),ch[root][1] = New(INF);//先加入正无穷和负无穷，便于之后操作(貌似不加也行)
 45     pushup(root);//因为INF > -INF,所以是右子树，
 46 }
 47 void Rotate(int &cur,int d) {
 48     int temp = ch[cur][d ^ 1];
 49     ch[cur][d ^ 1] = ch[temp][d];
 50     ch[temp][d] = cur;
 51     cur = temp;
 52     pushup(ch[cur][d]);
 53     pushup(cur);
 54     re;
 55 }
 56 void insert(int &id,int v) { //id依然是引用，在新建节点时可以体现
 57     if(!id) {
 58         id = ++tot;
 59         val[tot] = v;
 60         dat[tot] = rand();
 61         size[tot] = 1;
 62         cnt[tot] = 1;
 63         return;
 64     }
 65     if(v == val[id]) cnt[id]++;//若节点已存在，则副本数++;
 66     else { //要满足BST性质，小于插到左边，大于插到右边
 67         bool d = (v > val[id]);
 68         insert(ch[id][d],v);//递归实现
 69         if(dat[id] < dat[ch[id][d]]) Rotate(id,d ^ 1);//(参考一下图)与左节点交换右旋，与右节点交换左旋
 70     }
 71     pushup(id);//现在更新一下本节点的信息
 72 }
 73 void Del(int &id,int v) {
 74     if(!id)return ;//到这了发现查不到这个节点，该点不存在，直接返回
 75     if(v == val[id]) { //检索到了这个值
 76         if(cnt[id] > 1) {
 77             cnt[id]--,pushup(id);    //若副本不止一个，减去一个就好
 78             return ;
 79         }
 80         if(ch[id][0] || ch[id][1]) { //发现只有一个值，且有儿子节点,我们只能把值旋转到底部删除
 81             if(!ch[id][1] || dat[ch[id][0]] > dat[ch[id][1]]) { //当前点被移走之后，会有一个新的点补上来(左儿子或右儿子)，按照优先级，优先级大的补上来
 82                 Rotate(id,1),Del(ch[id][1],v);//我们会发现，右旋是与左儿子交换，当前点变成右节点；左旋则是与右儿子交换，当前点变为左节点
 83             } else Rotate(id,0),Del(ch[id][0],v);
 84             pushup(id);
 85         } else id = 0;//发现本节点是叶子节点，直接删除
 86         return ;//这个return对应的是检索到值de所有情况
 87     }
 88     v < val[id] ? Del(ch[id][0],v) : Del(ch[id][1],v);//继续BST性质
 89     pushup(id);
 90 }
 91 int get_rank(int id,int v) {
 92     if(!id)return 0;//若查询值不存在，返回
 93     if(v == val[id])return size[ch[id][0]] + 1;//查询到该值，由BST性质可知：该点左边值都比该点的值(查询值)小，故rank为左儿子大小 + 1
 94     else if(v < val[id])return get_rank(ch[id][0],v);
 95     //发现需查询的点在该点左边，往左边递归查询
 96     else return size[ch[id][0]] + cnt[id] + get_rank(ch[id][1],v);
 97     //若查询值大于该点值。说明询问点在当前点的右侧，且此点的值都小于查询值，所以要加上cnt[id]
 98 }
 99 int get_val(int id,int rank) {
100     if(!id) return INF;//一直向右找找不到，说明是正无穷
101     if(rank <= size[ch[id][0]]) return get_val(ch[id][0],rank);
102     //左边排名已经大于rank了，说明rank对应的值在左儿子那里
103     else if(rank <= size[ch[id][0]] + cnt[id]) return val[id];
104     //上一步排除了在左区间的情况，若是rank在左与中(目前节点)中，
105     //则直接返回目前节点(中区间)的值
106     else return get_val(ch[id][1],rank - size[ch[id][0]] - cnt[id]);
107     //剩下只能在右区间找了，rank减去左区间大小和中区间，继续递归
108 }
109 int get_pre(int v) {
110     int id = root,pre;
111     while(id) {
112         if(val[id] < v) pre = val[id],id = ch[id][1];
113         else id = ch[id][0];
114     }
115     return pre;
116 }
117 int get_next(int v) {
118     int id = root,next;
119     while(id) {
120         if(val[id] > v) next = val[id],id = ch[id][0];
121         else id = ch[id][1];
122     }
123     return next;
124 }
125 int main() {
126     build();
127     //不要忘记初始化
128     //[运行build()会连同root一并初始化，所以很重要]
129     na = read();
130     for(int i = 1; i <= na; i++) {
131         int cmd = read(),x = read();
132         if(cmd == 1)insert(root,x);//函数都写好了，注意：需要递归的函数都从根开始，不需要递归的函数直接查询
133         if(cmd == 2)Del(root,x);
134         if(cmd == 3)printf("%d\n",get_rank(root,x) - 1);//注意：因为初始化时插入了INF和-INF,所以查询排名时要减1(-INF不是第一小，是“第零小”)
135         if(cmd == 4)printf("%d\n",get_val(root,x + 1));//同理，用排名查询值得时候要查x + 1名，因为第一名(其实不是)是-INF
136         if(cmd == 5)printf("%d\n",get_pre(x));
137         if(cmd == 6)printf("%d\n",get_next(x));
138     }
139     return 0;
140 }