旋转Treap树

发布时间 2023-07-06 16:16:07作者: ZRJ0721
#include<bits/stdc++.h>
using namespace std;

const int M=1e6+10;
int cnt = 0; // cnt记录当前节点个数,最新一个节点即t[cnt]
int root = 0; // root是整棵树的树根,初始为空树所以root初始化=0
int n; // n表示操作次数

struct Node{
	int ls, rs; // 左右儿子
	int val, pri; // val:权值;pri:随机的优先级
	int size; // 当前结点为根的子树的结点数量,用于求第k大和rank
}t[M]; // tree[],存树

void newNode(int x){ // 初始化一个新结点
    // 新节点初始是叶子结点,之后在插入节点函数中调整位置
	cnt++; // 从t[1]开始存储结点,0表示空结点(若子节点为0即无子节点)
	t[cnt].size = 1; // 最初子树大小为1(即叶子结点自身)
	t[cnt].ls = t[cnt].rs = 0; // 叶子结点没有子结点
	t[cnt].val = x; // val: 赋予指定的键值
	t[cnt].pri = rand(); // pri:随机产生一个优先级,rand()是随机数函数
}

void update(int u){ // 更新以u为根的子树的size
	t[u].size = t[t[u].ls].size + t[t[u].rs].size+1;
}

void rotate(int &a, bool d){ // 当前位置的节点,当前操作
    // 具体地,d=true则左旋,d=false则右旋
    // a表示当前位置的节点:在旋转中,当前位置的节点会变
    // int& a表示a为一个int值的地址
    // 简单来说,即使得修改a的同时也修改函数外的值
    // 使用返回值可以达到同样的效果
    int b; // 节点B,同时是新插入的节点
    if(d == true){ // 左旋
        b = t[a].rs; // B是A的右儿子(right son)
        t[a].rs = t[b].ls;
        // 这里是解决儿子被顶替的情况,即上文的节点C
        t[b].ls = a; // A变成B的左儿子(left son)
    }
    else{ // 右旋,操作差不多,方向反过来
        b = t[a].ls; // B是A的左儿子
        t[a].ls = t[b].rs;
        t[b].rs = a; // A变成B的右儿子
    }
    t[b].size = t[a].size; // 继承子树大小
    update(a); // 重新计算A点的size,此处略
    a = b; // 此处a不再代表A,而代表子树的“根”
    // B节点代替A成为子树的“根”
} 

void insert(int &u, int x){ // 当前位置的节点 插入的权值
    // u和旋转部分函数的a同理
	if(u == 0){ // 假如找到了空位(空叶子节点)
        newNode(x); // 初始化权值为x的新节点,此处略
        u = cnt; // cnt是节点个数,同时也是新节点的编号
        return; // 直接返回
    }
    // 则新节点在当前位置的子树中
	t[u].size++; // 子树大小累加
    // 注:val是权值,pri是优先级
	if(x >= t[u].val) //如果新权值大于当前点
        insert(t[u].rs,x); //递归右子树找空叶子,直到插入
	else // 如果新权值小于当前点
        insert(t[u].ls,x); //同上,递归左子树
	if(t[u].ls!=0 && t[u].pri>t[t[u].ls].pri)
        rotate(u, 0); // 判断是否需要旋转
	if(t[u].rs!=0 && t[u].pri>t[t[u].rs].pri)
        rotate(u, 1);
	update(u); // 重新计算当前位置的size,此处略
}

void del(int &u, int x){ // 当前遍历到的节点 要删除的节点
    // u和旋转部分函数的a同理
    // 遍历到u意味着:要删除的节点必然在u的子树中
	t[u].size--; // 删除目标节点后u子树中少了一个节点
	if(t[u].val == x){ // 如果当前的节点就是要删除的节点
		if(t[u].ls == 0 && t[u].rs == 0){ // 如果u无依无靠、孤苦伶仃
            u = 0; // 直接把它删除掉
            return; // 返回
        } // 注意一下,这些if中只要返回了就不会进行之后的判断和计算
        // 相当于if else if结构
		if(t[u].ls == 0 || t[u].rs == 0){ // 如果u只有一个儿子
            u = t[u].ls + t[u].rs; // 把u修改成儿子,相当于移动到儿子节点上
            // 这一步是为了函数中最后一行,更新儿子的size值
            // 可能用返回值会更好理解?
            return;
        }
		if(t[t[u].ls].pri < t[t[u].rs].pri){
			rotate(u, 0);
			del(t[u].rs, x);
			return;
		}
		else{
			rotate(u, 1);
			del(t[u].ls, x);
			return;
		}
	}
	if(t[u].val >= x) // 经典的二分查找部分,略
        del(t[u].ls,x); // 往左子树找
	else del(t[u].rs,x); // 在右子树找
    // 此时必然已经删除完毕,"u"是被删节点的儿子
	update(u); // 更新size值
}

int Rank(int u, int x){ // 排名(从小到大排序,x为第几个)
	if(u == 0) return 0;
	if(x > t[u].val)
        return t[t[u].ls].size + 1 + Rank(t[u].rs, x);
	return Rank(t[u].ls, x);
}

int kth(int u,int k){ // 求第k大数
    // 注意:事实上,函数的k是表示求“以u为根的子树中的第k大数”
    // k会在函数传递的过程中改变(这是计划的一部分)
	if(k == t[t[u].ls].size + 1) // 如果这个数就是子树中第k大的数
        return t[u].val; // u刚好就是要查找的数
	else // 如果不是要找的数,就继续二分查找
        if(k > t[t[u].ls].size+1) // 如果第k大在右子树里
            return kth(t[u].rs, k-t[t[u].ls].size-1);
    		// 右子树(为了排除左子树中较小的数,要修改k的值)
		else return kth(t[u].ls, k); // 左子树(因为左子树就是u子树里最小的一部分,所以不用修改k)
}

int precursor(int u, int x){ // 求x的前驱(略)
	if(u == 0) return 0;
	if(t[u].val >= x)
        return precursor(t[u].ls, x);
	int tmp = precursor(t[u].rs,x);
	if(tmp == 0) return t[u].val;
	return tmp;
}

int successor(int u,int x){ // 求x的后继(略)
	if(u == 0) return 0;
	if(t[u].val <= x) return successor(t[u].rs,x);
	int tmp = successor(t[u].ls,x);
	if(tmp == 0) return t[u].val;
	return tmp;
}

signed main(){
	scanf("%d",&n); // 输入n
	while(n--){ // 循环n次
        int opt, x; // opt表示操作类型,x为对应的值
        scanf("%d%d", &opt, &x);
        switch (opt){ // 避免大量if的出现
            case 1: insert(root, x); break; // 插入新节点
            case 2: del(root,x); break; // 删除节点
            case 3: printf("%d\n", Rank(root, x)+1); break; // 输出排名
            case 4: printf("%d\n", kth(root, x)); break; // 输出第x大数
            case 5: printf("%d\n", precursor(root, x)); break; // 输出x前驱
            case 6: printf("%d\n", successor(root, x)); break; // 输出x后继
        }
	}
	return 0;
}