splay模板(封装)

Janokay · 2022-08-28 19:16:33 · 个人记录

#pragma GCC optimize(2)
#include <iostream>

using namespace std;

const int N = 1e5 + 10;

class Splay_Tree
{
    #define fa(x) tr[x].fa
    #define ls(x) tr[x].ch[0]
    #define rs(x) tr[x].ch[1]
    #define root tr[0].ch[1]
    public:
        int idx;
        struct Splay
        {
            int ch[2];
            int fa, cnt, val, size;
        }tr[N];

        int ident(int x) {return tr[fa(x)].ch[1] == x;}

        void update(int x) {tr[x].size = tr[ls(x)].size + tr[rs(x)].size + tr[x].cnt;}

        int set_node(int v, int f)
        {
            this->idx ++ ;
            tr[this->idx].val = v; tr[this->idx].fa = f;
            tr[this->idx].cnt = tr[this->idx].size = 1;
            return this->idx;
        }

        void connect(int x, int y, int how)
        {
            tr[y].ch[how] = x;
            tr[x].fa = y;
        }

        void rotate(int x)
        {
            int Y = tr[x].fa, R = tr[Y].fa;
            int Yson = ident(x), Rson = ident(Y);
            connect(tr[x].ch[Yson ^ 1], Y, Yson);
            connect(Y, x, Yson ^ 1);
            connect(x, R, Rson);
            update(Y);
            update(x);
        }

        void splay(int x, int to)
        {
            to = tr[to].fa;
            while (fa(x) ^ to)
            {
                int y = fa(x);
                if (fa(y) == to)
                    rotate(x);
                else if (ident(x) == ident(y))
                    rotate(y), rotate(x);
                else 
                    rotate(x), rotate(x);
            }
        }

        void insert(int x)
        {
            int now = root;
            if (!now) 
                root = set_node(x, 0);
            else
            {
                while (1)
                {
                    tr[now].size ++ ;
                    if (tr[now].val == x)
                    {
                        tr[now].cnt ++ ;
                        splay(now, root);
                        return ;
                    }
                    int nxt = x >= tr[now].val;
                    if (!tr[now].ch[nxt])
                    {
                        int p = set_node(x, now);
                        tr[now].ch[nxt] = p;
                        splay(p, root);
                        return ;
                    }
                    now = tr[now].ch[nxt];
                }
            }
        }

        int find(int x)
        {
            int now = root;
            while (1)
            {
                if (!now) return 0;
                if (tr[now].val == x)
                {
                    splay(now, root);
                    return now;
                }
                int nxt = x >= tr[now].val;
                now = tr[now].ch[nxt];
            }
        }

        void delet(int x)
        {
            int pos = find(x);
            if (!pos) return ;
            if (tr[pos].cnt > 1)
            {
                tr[pos].cnt -- ; tr[pos].size -- ;
                return ;
            }
            else
            {
                if (!tr[pos].ch[0] && !tr[pos].ch[1])
                {
                    root = 0;
                    return ;
                }
                else if (!tr[pos].ch[0])
                {
                    root = tr[pos].ch[1];
                    tr[root].fa = 0;
                    return ;
                }
                else
                {
                    int left = tr[pos].ch[0];
                    while (rs(left)) left = rs(left);
                    splay(left, tr[pos].ch[0]);
                    connect(tr[pos].ch[1], left, 1);
                    connect(left, 0, 1);
                    update(left);
                }
            }
        }

        int rank(int x)
        {
            return tr[ls(find(x))].size + 1;
        }

        int value(int x)
        {
            int now = root;
            while (1)
            {
                int tem_num = tr[now].size - tr[rs(now)].size;
                if (tr[ls(now)].size < x && x <= tem_num)
                {
                    splay(now, root);
                    return tr[now].val;
                }
                if (x < tem_num) now = ls(now);
                else now = rs(now), x -= tem_num;
            }
        }

        int prev(int x)
        {
            int now = root, ans = -0x3f3f3f3f;
            while (now)
            {
                if (tr[now].val < x)
                    ans = max(ans, tr[now].val);
                int nxt = x <= tr[now].val ? 0 : 1;
                now = tr[now].ch[nxt];
            }
            return ans;
        }

        int nxt(int x)
        {
            int now = root, ans = 0x3f3f3f3f;
            while (now)
            {
                if (tr[now].val > x)
                    ans = min(ans, tr[now].val);
                int nxt = x < tr[now].val ? 0 : 1;
                now = tr[now].ch[nxt];
            }
            return ans;
        }

}tree;

int n;
int main(void)
{
    scanf("%d", &n);
    while (n -- )
    {
        int opt, x;
        scanf("%d%d", &opt, &x);
        if (opt == 1) tree.insert(x);
        else if (opt == 2) tree.delet(x);
        else if (opt == 3) printf("%d\n", tree.rank(x)); 
        else if (opt == 4) printf("%d\n", tree.value(x));
        else if (opt == 5) printf("%d\n", tree.prev(x));
        else printf("%d\n", tree.nxt(x));
    }
    return 0;
}