萌新求调平衡树板子
查看原帖
萌新求调平衡树板子
1069533
sto_clx_orz楼主2024/7/9 09:20

RT

#include<bits/stdc++.h>
using namespace std;

namespace clx
{
    class AVL
    {
        struct point
        {
            int key,value;
            int height,size,cnt;
            AVL*tree;
            point *child[2],*father;
            point(int k,int v,AVL*t):tree(t),value(v),key(k),height(1),size(0),cnt(0),child({nullptr,nullptr}){}
        }*root;
        int get_height(point*u)
        {return u==nullptr?0:u->height;}
        int get_size(point*u)
        {return u==nullptr?0:u->size;}
        void pushup(point*u)
        {
            if(u==nullptr)
                return;
            u->size=get_size(u->child[0])+get_size(u->child[1])+u->cnt;
            u->height=max(get_height(u->child[0]),get_height(u->child[1]))+1;
        }
        int get(point*u)
        {return u==u->father->child[1];}
        void rotate(point*u)
        {
            if(u==nullptr||u==root)
                return;
            if(u->father==root)
                root=u;
            int b=get(u);
            point*father=u->father;
            father->child[b]=u->child[!b];
            if(u->child[!b]!=nullptr)
                u->child[!b]->father=father;
            u->child[!b]=father;
            u->father=father->father;
            if(u->father==nullptr)
                root=u;
            else 
                father->father->child[get(father)]=u;
            father->father=u;
            pushup(father),pushup(u);
            if(u!=root)
                pushup(u->father);
        }
        point*adjustment(point*u)
        {
            if(abs(get_height(u->child[0])-get_height(u->child[1]))<2)
                return pushup(u),u;
            point*v,*x;
            int b;
            if(get_height(u->child[0])-get_height(u->child[1])==2)
                v=u->child[0],b=0;
            else
                v=u->child[1],b=1;
            if(get_height(v->child[b])>=get_height(v->child[!b]))
                rotate(x=v);
            else
                rotate(v->child[!b]),rotate(x=v->child[!b]);
            return x;
        }
        point*query_by_rank(int rnk,point*now)
        {
            while(now!=nullptr)
            {
                if(rnk>=get_size(now->child[0])+1&&rnk<=get_size(now->child[0])+now->cnt)
                    return now;
                if(rnk>get_size(now->child[0])+now->cnt)
                    rnk-=get_size(now->child[0])+now->cnt,
                    now=now->child[1];
                else
                    now=now->child[0];
            }
            return now;
        }
        void Destroy(point*u)
        {
            if(u->child[0]!=nullptr)
                Destroy(u->child[0]),
                delete u->child[0];
            if(u->child[1]!=nullptr)
                Destroy(u->child[1]),
                delete u->child[1];
        }
    public:
        AVL():root(nullptr){}
        ~AVL()
        {
            if(root!=nullptr)
                Destroy(root),
                delete root;
        }
        void insert(int key,int value)
        {
            point*now=root,*New=new point(key,value,this);
            if(root==nullptr)
                return New->cnt++,New->size++,root=New,void();
            int b=key>now->key;
            while(now->child[b]!=nullptr&&now->key!=key)
                now=now->child[b],b=key>now->key;
            if(now->key!=key)
                New->father=now,
                now->child[b]=New,
                ++New->cnt,
                ++New->size;
            else 
            {
                delete New;
                ++now->cnt,
                ++now->size,
                now=now->father;
            }
            while(now!=nullptr)
            {
                now=adjustment(now);
                if(now!=nullptr)
                    now=now->father;
            }
        }
        int query_rank(int key)
        {
            point*now=root;
            int rank=1;
            while(now!=nullptr)
                if(key>now->key)
                    rank+=get_size(now->child[0])+now->cnt,
                    now=now->child[1];
                else
                    now=now->child[0];
            return rank;
        }
        class iterator
        {
            point*now;
        public:
            iterator(point*p):now(p){}
            iterator&operator++()
            {*this=now->tree->query_by_rank(now->tree->query_rank(now->key+1),now->tree->root);return*this;}
            iterator&operator--()
            {*this=now->tree->query_by_rank(now->tree->query_rank(now->key)-1,now->tree->root);return*this;}
            pair<int,int>operator*(){return {now->key,now->value};}
            bool operator!=(const iterator&it)const{return now!=it.now;}
            bool operator==(const iterator&it)const{return now==it.now;}
            friend AVL;
        };
        iterator query_by_rank(int rnk)
        {return query_by_rank(rnk,root);}
        void erase(int key)
        {
            point*now=root;
            while(now!=nullptr)
            {
                if(now->key==key)break;
                now=now->child[key>now->key];
            }
            if(now==nullptr)
                return;
            if(now->cnt>1)
            {
                --now->cnt;
                while(now!=nullptr)
                    pushup(now),now=now->father;
                return;
            }
            if(now->child[0]!=nullptr&&now->child[1]!=nullptr)
            {
                point*nxt=query_by_rank(1,now->child[1]);
                swap(now->key,nxt->key);
                swap(now->value,nxt->value);
                swap(now->cnt,nxt->cnt);
                now=nxt;
            }
            if(now->father==nullptr)
                if(now->child[0]!=nullptr)
                    root=now->child[0];
                else 
                    root=now->child[1];
            else 
                if(now->child[0]!=nullptr)
                    now->father->child[get(now)]=now->child[0],
                    now->child[0]->father=now->father;
                else 
                {
                    now->father->child[get(now)]=now->child[1];
                    if(now->child[1]!=nullptr)
                        now->child[1]->father=now->father;
                }
            point*father=now->father;
            delete now;
            while(father!=nullptr)
            {
                father=adjustment(father);
                if(father!=nullptr)
                    father=father->father;
            }
        }
        iterator begin()
        {return query_by_rank(1);}
        iterator end()
        {return nullptr;}
    };
}
using namespace clx;

AVL avl;
int n;

int main()
{
    ios::sync_with_stdio(false),cin.tie(0);
    cin>>n;
    while(n--)
    {
        int opt,x;
        cin>>opt>>x;
        switch(opt)
        {
            case 1:
                avl.insert(x,x);
                break;
            case 2:
                avl.erase(x);
                break;
            case 3:
                cout<<avl.query_rank(x)<<'\n';
                break;
            case 4:
                cout<<(*avl.query_by_rank(x)).first<<'\n';
                break;
            case 5:
                cout<<(*--avl.query_by_rank(avl.query_rank(x))).first<<'\n';
                break;
            case 6:
                cout<<(*avl.query_by_rank(avl.query_rank(x+1))).first<<'\n';
        }
    }
}
2024/7/9 09:20
加载中...