相比于之前的普通平衡树进行左旋右旋来比,splay的适用性更高,使用更广泛。
核心函数rotate、splay函数,其它的根据需要进行修改。
int n, m;
struct Node {
int s[2], p, v, cnt; // 左右儿子、父节点、值、出现数量
int size, flag; // 子树大小、懒标记
void init(int _v, int _p) { // 初始化函数
v = _v, p = _p;
cnt = size = 1;
}
} tr[N];
int root, idx;// 根节点、分配节点序号
void pushup(int u) { // 向上更新传递,与线段树一样
tr[u].size = tr[tr[u].s[0]].size + tr[tr[u].s[1]].size + tr[u].cnt;
}
void pushdown(int x) { // 向下传递更新 ,与线段树一样
if(tr[x].flag) {
swap(tr[x].s[0], tr[x].s[1]);
tr[tr[x].s[0]].flag ^= 1;
tr[tr[x].s[1]].flag ^= 1;
tr[x].flag = 0;
}
}
void rotate(int x) { // 核心函数
int y = tr[x].p, z = tr[y].p;
int k = tr[y].s[1] == x;
tr[z].s[tr[z].s[1] == y] = x, tr[x].p = z;
tr[y].s[k] = tr[x].s[k ^ 1], tr[tr[x].s[k ^ 1]].p = y;
tr[x].s[k ^ 1] = y, tr[y].p = x;
pushup(y), pushup(x);
}
void splay(int x, int k) { // 将x节点旋转到k节点下
while(tr[x].p != k) { //
int y = tr[x].p; // x节点的父节点
int z = tr[y].p; // x节点的父节点的父节点
if(z != k) // 向上旋转
if((tr[y].s[1] == x) != (tr[z].s[1] == y)) rotate(x); // 转一次x
else rotate(y); // 转一次y
rotate(x); // 转一次x
}
if(!k) root = x; // 更新root节点
}
void upper(int v) { // 将v值节点转到根节点
int u = root; // 根节点
while(tr[u].s[v > tr[u].v] && tr[u].v != v) // 存在则找到v值节点,不存在则找到v值节点的前驱或者后继节点
u = tr[u].s[v > tr[u].v]; // 向下寻找
splay(u, 0); // 将u节点旋转到跟节点
}
int get_prev(int v) { // 获取v值的前驱节点
upper(v); // 将v值节点转到根节点
if(tr[root].v < v) return root; // 若是该值在树中不存在,根节点就是v的前驱或者后继节点
int u = tr[root].s[0]; // 前驱节点在左子树的最右边
while(tr[u].s[1]) u = tr[u].s[1]; // 找到最右边的一个节点
return u;
}
int get_next(int v) { // 获取某值的后继节点
upper(v); // 将v值节点转到根节点
if(tr[root].v > v) return root; // 若是该值在树中不存在,根节点就是v的前驱或者后继节点
int u = tr[root].s[1]; // 后继节点在右子树的最左边
while(tr[u].s[0]) u = tr[u].s[0]; // 找到最左的节点,就是最小的节点
return u; // 返回节点
}
int get_rank_by_key(int v) { // key值在当前树中的排名
upper(v); //
if(tr[root].v >= v)
return tr[tr[root].s[0]].size + 1;
return tr[tr[root].s[0]].size + tr[root].cnt + 1;
}
int get_key_by_rank(int k) { // 获取树中排名为k的值
int u = root; // 根节点
while(tr[u].size >= k) { // 保证当前子树中有解
if(tr[tr[u].s[0]].size >= k) u = tr[u].s[0]; // 在左子树中
else if(tr[tr[u].s[0]].size + tr[u].cnt >= k) return splay(u, 0), tr[u].v; // 在当前节点
else k -= tr[tr[u].s[0]].size + tr[u].cnt, u = tr[u].s[1]; // 在右子树,需要更新k值,减去左子树以及当前节点值的数量
}
return -1;
}
void insert(int v) { // 在二叉树中插入一个值
int u = root, p = 0; // p维护为当前节点的父节点
while(u && tr[u].v != v) // 没找到则一直向下寻找
p = u, u = tr[u].s[v > tr[u].v]; // 更新父节点,更新当前节点
if(u) tr[u].cnt ++; // v值的节点已经存在则直接加一即可
else { // 不存在则创建节点
u = ++ idx; // 分配节点序号
if(p) tr[p].s[v > tr[p].v] = u; // 将父节点也就是前驱节点指向当前节点
tr[u].init(v, p); // 初始化当前节点的值、父节点信息
}
splay(u, 0); // 将u节点旋转到根节点下
}
void remove(int v) { // 删除一个值为v的节点
int prev = get_prev(v), nex = get_next(v); // 获取该节点的前驱以及后继节点。
splay(prev, 0), splay(nex, prev); // 将前继节点旋转到根节点,将后继节点旋转到前驱节点下面也就是根节点下面
int w = tr[nex].s[0]; // 后继节点的左子树就是v的节点
if(tr[w].cnt > 1) tr[w].cnt --, splay(w, 0); // 该节点的v不止存在一个,减一,w节点旋转到根节点
else tr[nex].s[0] = 0, splay(nex, 0); // 唯一,那么直接把后继节点的左子树指向空也就是0即可
}
void output(int u) { // 中序遍历输出二叉树
// pushdown(u);
int l = tr[u].s[0], r = tr[u].s[1]; // 左右儿子
if(l) output(l); // 递归左儿子
if(tr[u].v >= 1 && tr[u].v <= n) cout << tr[u].v << " "; // 输出当前子树的根
if(r) output(r); // 递归右儿子
}
inline void sovle() {
cin >> n;
insert(-INF), insert(INF);
// 插入两个哨兵,无穷小以及无穷大 使得在查询某数不存在的是时候不会产生越界
while(n --) {
int a, b;
cin >> a >> b;
if(a == 1) insert(b); // 插入一个值
if(a == 2) remove(b); // 插入一个值
if(a == 3) cout << get_rank_by_key(b) - 1 << endl; // 真实排名减一 因为前面多了一个哨兵
if(a == 4) cout << get_key_by_rank(b + 1) << endl; // 真实排名加一 因为哨兵
if(a == 5) cout << tr[get_prev(b)].v << endl;
if(a == 6) cout << tr[get_next(b)].v << endl;
}
}