SPLAY平衡树【模板+例题】TOJ6591/luogu3369

您需要写一种数据结构（可参考题目标题），来维护一些数，其中需要提供以下操作：

1、插入 x 数；

2 、删除 x 数（若有多个相同的数，因只删除一个）；

3、查询 x 数的排名（若有多个相同的数，因输出最小的排名）；

4、查询排名为 x 的数；

5、求 x 的前趋（前趋定义为小于 x，且最大的数）；

6、求 x 的后继（后继定义为大于 x，且最小的数）。

链接：TOJ6591

#include <bits/stdc++.h>
using namespace std;
#define inf 0x3f3f3f3f
typedef long long LL;
typedef unsigned long long ULL;
const int N = 1e5+10;
int mod = 1e9+7;
int P = 13331;

int n, root, idx;//root为根节点，iex为树节点个数
struct node
{
    int s[2];//s[0]代表左儿子，s[1]代表右儿子
    int p;//节点的父亲
    int v;//节点的权值
    int size;//以该点为根的子树大小
    int cnt;//该权值的数目
    void init(int p1, int v1) //节点初始化
    {
        p = p1;
        v = v1;
        size = cnt = 1;
    }
}tr[N];
void pushup(int x)//旋转之后需要重新计算子树大小
{
    tr[x].size = tr[tr[x].s[0]].size+tr[tr[x].s[1]].size+tr[x].cnt;//加俩儿子大小和与自己大小相同的个数
}
void rotate(int x)//旋转子树
{
    int y = tr[x].p, z = tr[y].p;//y为x节点的父亲，z为y节点的父亲
    int k = tr[y].s[1] == x;//判断x为y的左儿子还是右儿子
    tr[y].s[k] = tr[x].s[1^k];//y的左/右儿子替换为x的右/左儿子
    tr[tr[x].s[1^k]].p = y;//x的右/左儿子替换为y
    tr[x].s[k^1] = y;//x和y父子交换左换右，右换左
    tr[y].p = x;//y的父亲变为x
    tr[z].s[tr[z].s[1] == y] = x;//z的左/右儿子变为x
    tr[x].p = z;//x的父亲变为z
    pushup(x), pushup(y);//更新x和y的子树大小
}
void splay(int x,int k)//将x节点转到k位置下面
{
    //k == 0时，将x节点转到根节点
    while(tr[x].p != k)
    {
        int y = tr[x].p, z = tr[y].p;
        //如果y是根直接旋转，否则需要再转
        if (z != k)
        {
            //判断x,y,z的形状是直线型还是折线型
            if ((tr[y].s[0] == x)^(tr[z].s[0] == y))//折线型转x
                rotate(x);
            else//直线型转y
                rotate(y);
        }
        rotate(x);
        // cout << "====" << endl;
    }

    if (k == 0) root = x;
}
void find(int v)//查找权值为v的值，并将其转为根节点
{
    int x = root;
    while (tr[x].s[v>tr[x].v]&&v != tr[x].v)//存在这个节点，相同返回，否则根据大小判断转到左/右儿子
    {
        x = tr[x].s[v>tr[x].v];
    }
    splay(x,0);
}
int get_pre(int v)//求v的前驱
{
    find(v);
    int x = root;
    if (tr[x].v < v) return x;//不存在v，而且根权值小于v，此时根节点为前驱直接返回
    x = tr[x].s[0];
    while (tr[x].s[1]) x = tr[x].s[1];
    splay(x, 0);
    return x;
}
int get_suc(int v)//求v的后继
{
    find(v);
    int x = root;
    if (tr[x].v > v) return x;//不存在v，而且根权值大于v，此时根节点为后驱直接返回
    x = tr[x].s[1];
    while (tr[x].s[0]) x = tr[x].s[0];
    splay(x, 0);
    return x;
}
void del(int v)//将v删除
{
    int pre = get_pre(v);
    int suc = get_suc(v);
    splay(pre,0), splay(suc,pre);//将v转为叶子节点方便删除
    int del = tr[suc].s[0];
    if (tr[del].cnt > 1)//如果个数不止一个删去一个
        tr[del].cnt--, splay(del,0);
    else//否则直接删去
        tr[suc].s[0] = 0, splay(suc,0);
}
int get_rank(int v)//查找v的排名
{
    find(v);
    return tr[tr[root].s[0]].size;
}
int get_val(int k)//查询排名为k的树
{
    int x = root;
    while(1)
    {
        int y = tr[x].s[0];
        if (tr[y].size+tr[x].cnt < k)//左子树加根个数小于k，证明第k个数在右子树
        {
            k -= tr[y].size+tr[x].cnt;
            x = tr[x].s[1];
        }
        else//如果左子树个数大于k，在左子树，否则在当前根上；
        {
            if (tr[y].size >= k) x = tr[x].s[0];
            else break;
        }
    }
    splay(x,0);
    return tr[x].v;
}
void insert(int v) //插入大小为v的节点
{
    int x = root, p = 0;
    while (x&&tr[x].v != v)//找应该插入的位置
    {
        p = x;
        x = tr[x].s[v>tr[x].v];
    }
    if (x) tr[x].cnt++;//如果存在个数加1
    else
    {
        x = ++idx;
        tr[p].s[v>tr[p].v] = x;
        tr[x].init(p,v);
    }
    splay(x,0);
}
void solve()
{
    insert(-1e9);//左终点
    insert(1e9);//右终点
    cin >> n;
    while(n--)
    {
        int op, x;
        cin >> op >> x;
        if (op == 1) insert(x);
        if (op == 2) del(x);
        if (op == 3) cout << get_rank(x) << endl;
        if (op == 4) cout << get_val(x+1) << endl;
        if (op == 5) cout << tr[get_pre(x)].v << endl;
        if (op == 6) cout << tr[get_suc(x)].v << endl;
    }
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0), cout.tie(0);
    int t = 1;
    // cin >> t;
    while (t--)
    {
        solve();
    }
	return 0;
}