C++写红黑树 翻仓库发现的大宝贝

C++写红黑树

洛谷P3369

/*
    https://www.luogu/problem/P3369
*/
#include<iostream>
#include<cstring>
#include<cstdio>
using namespace std;
#define INF 0x3f3f3f3f
#define maxn 105
#define minn -105
#define ll long long int
#define islc(x) ((x)&&(x)->fa->lc==(x))
#define isrc(x) ((x)&&(x)->fa->rc==(x))
#define bro(x) ((x)->fa->lc==(x)?((x)->fa->rc):((x)->fa->lc))

struct node
{
    node* fa; //父节点
    node* lc; //左孩子
    node* rc; //右孩子
    int val; //值
    bool RBc; //颜色，黑为0，红为1
    int num; //以this为根的子树的节点数
    node(int v,bool c,node* f=NULL,int s=1,node* l=NULL,node* r=NULL)
    {
        val=v;RBc=c;num=s;
        fa=f;lc=l;rc=r;
    }
    node(){ }
    void maintain() //维护当前节点数
    {
        num=1;
        if(lc)num+=lc->num;
        if(rc)num+=rc->num;
    }
    node* left_node() //寻找直接前驱，可能不在this下方
    {
        node* ptr=this;
        if(!lc) //无左子树，只能向上寻找
        {
            while(ptr->fa&&islc(ptr))ptr=ptr->fa; //只要当前ptr是左子节点，父节点肯定不比它小
            ptr=ptr->fa; //如果this是最小值，会返回NULL
        }
        else
        {
            ptr=ptr->lc;
            while(ptr->rc)ptr=ptr->rc;
        }
        return ptr;
    }
    node* right_node() //寻找直接后驱，可能不在this下方
    {
        node* ptr=this;
        if(!rc)
        {
            while(ptr->fa&&isrc(ptr))ptr=ptr->fa;
            ptr=ptr->fa;
        }
        else
        {
            ptr=ptr->rc;
            while(ptr->lc)ptr=ptr->lc;
        }
        return ptr;
    }
    node* succ() //后面删除过程中，要查找后继交换节点，即右支中值域最小节点
    {
        node* ptr=rc;
        while(ptr->lc)
        {
            ptr->num--; //注意这里num--
            ptr=ptr->lc;
        }
        return ptr;
    }
};

node* root(NULL); //树根

node* find_insert_pos(int val,int op) //查找插入的位置,tmp记录插入节点
{
    node* tmp=NULL;
    node* ptr=root;
    while(ptr)
    {
        tmp=ptr;
        ptr->num+=op; //在一些查找中顺带进行修改
        if(ptr->val>val)ptr=ptr->lc; //相等也往右走
        else ptr=ptr->rc;
    }
    return tmp;
}

node* find_node(int val,int op) //查找，同值节点中随便返回一个
{
    node* ptr=root;
    while(ptr&&(ptr->val!=val)) //与上面的差别在这，同值直接结束
    {
        ptr->num+=op;
        if(ptr->val>val)ptr=ptr->lc;
        else ptr=ptr->rc;
    }
    return ptr;
}

void left_revolve(node* cur_gfa,node* new_gfa) //左旋,一定要非空
{
    node* ptr=new_gfa->lc;
    if(cur_gfa==root)new_gfa->fa=NULL,root=new_gfa;
    else
    {
        if(islc(cur_gfa))cur_gfa->fa->lc=new_gfa; //维护此旋转子树与上方父节点的关系
        else cur_gfa->fa->rc=new_gfa;
        new_gfa->fa=cur_gfa->fa;
    }
    new_gfa->lc=cur_gfa,cur_gfa->fa=new_gfa;
    if(ptr)ptr->fa=cur_gfa;
    cur_gfa->rc=ptr; //可能为空
    cur_gfa->maintain(),new_gfa->maintain(); //维护当前节点数量
}

void right_revolve(node* cur_gfa,node* new_gfa) //右旋，一定要非空
{
    node* ptr=new_gfa->rc;
    if(cur_gfa==root)new_gfa->fa=NULL,root=new_gfa;
    else
    {
        if(islc(cur_gfa))cur_gfa->fa->lc=new_gfa;
        else cur_gfa->fa->rc=new_gfa;
        new_gfa->fa=cur_gfa->fa;
    }
    new_gfa->rc=cur_gfa,cur_gfa->fa=new_gfa;
    if(ptr)ptr->fa=cur_gfa;
    cur_gfa->lc=ptr;
    cur_gfa->maintain(),new_gfa->maintain();
}

/*插入完成后进行颜色修正,ptr颜色必为红，
ptr可能是新插入的叶子节点，也可能进行过递归操作的子树的根(所以黑高是不确定的）*/
void solve_double_red(node* ptr)
{
    while(!ptr->fa||ptr->fa->RBc) //向上迭代,出现双红就进行修正
    {
        if(ptr==root) //特判
        {
            ptr->RBc=0;
            return;
        }
        node* faptr=ptr->fa;
        node* uncptr=bro(faptr);
        node* gfaptr=faptr->fa;
        if(uncptr&&uncptr->RBc) //叔叔节点是红色，改色，向上递归进行修正
        {
            gfaptr->RBc=1;
            faptr->RBc=uncptr->RBc=0;
            ptr=gfaptr;
            continue;
        }
        //以下叔叔节点不存在或为黑色
        if(isrc(faptr)) //注意祖父关系与父子关系不同侧，其情况不同
        {
            if(islc(ptr))right_revolve(faptr,ptr); //非同侧,双红进行单旋，使之同侧，并且不改变向下的颜色性质
            faptr=gfaptr->rc; //faptr更新成当前真正的父节点
            left_revolve(gfaptr,faptr);
            faptr->RBc=0,gfaptr->RBc=1;
            break;
        }
        else
        {
            if(isrc(ptr))left_revolve(faptr,ptr);
            faptr=gfaptr->lc;
            right_revolve(gfaptr,faptr);
            faptr->RBc=0,gfaptr->RBc=1;
            break;
        }
    }
}

void insert_node(int val) //插入新点
{
    node* tmp=find_insert_pos(val,1); //返回插入位置的父节点
    if(!tmp) //插入节点的父节点
    {
        root=new node(val,0);
        return;
    }
    node* ptr=new node(val,1,tmp);
    if(tmp->val<=val)tmp->rc=ptr; //注意等于这时往右支走
    else tmp->lc=ptr;
    solve_double_red(ptr); //颜色修正,ptr一定为红
}

node* find_lower(int val) //查找最后一个比特定值严格小于的数
{
    node* ptr=root;
    node* tmp=NULL; //父节点
    while(ptr)
    {
        tmp=ptr;
        if(ptr->val<val)ptr=ptr->rc; //注意val相等的点插在右支
        else ptr=ptr->lc; //ptr->val==val往左走
    }
    if(!tmp)return NULL;
    if(tmp->val<val)return tmp;
    return tmp->left_node();
}

node* find_higher(int val) //查找第一个比特定值严格大于的数
{
    node* ptr=root;
    node* tmp=NULL;
    while(ptr)
    {
        tmp=ptr;
        if(ptr->val>val)ptr=ptr->lc; //相等默认往右走
        else ptr=ptr->rc;
    }
    if(!tmp)return NULL;
    if(tmp->val>val)return tmp;
    return tmp->right_node();
}

int get_k_rank(int val,node* ptr)
{
    if(!ptr)return 1;
    if(ptr->val>=val)return get_k_rank(val,ptr->lc);
    int ans=(ptr->lc)?ptr->lc->num:0; //这时查询值一定严格大于左支任意一个节点
    ans++;
    return ans+get_k_rank(val,ptr->rc); //右支可能还有比他小的
}

int query_rank_k(int rak,node* ptr) //ptr非空
{
    if(!ptr)return -1;
    if((ptr->lc)==NULL)
    {
        if(rak==1)return ptr->val;
        return query_rank_k(rak-1,ptr->rc);
    }
    if(ptr->lc->num>=rak)return query_rank_k(rak,ptr->lc);
    rak-=ptr->lc->num;
    if(rak==1)return ptr->val;
    else return query_rank_k(rak-1,ptr->rc);
}

void solve_double_black(node* ptr) //黑色修正,思路是让当前删除位置的一支黑高+1
{
    while(ptr!=root&&!ptr->RBc) //兄弟节点必定非空
    {
        node* broptr=bro(ptr);
        node* faptr=ptr->fa;
        if(broptr->RBc) //如果兄弟节点为红色，进行单旋
        {
            if(islc(broptr))right_revolve(faptr,broptr);
            else left_revolve(faptr,broptr);
            faptr->RBc=1,broptr->RBc=0;
            broptr=bro(ptr); //更新现在的bro
        }
        //注意接下来的兄弟节点一定非空，故它只能为黑色
        node* insdc; //兄弟内向儿子
        node* outsdc; //兄弟外向儿子
        if(islc(broptr))insdc=broptr->rc,outsdc=broptr->lc;
        else insdc=broptr->lc,outsdc=broptr->rc;
        if(!(insdc&&insdc->RBc)&&!(outsdc&&outsdc->RBc)) //兄弟两个孩子不存在或为黑
        {
            broptr->RBc=1; //兄弟为红
            if(!faptr->RBc) //父为黑，向上迭代修改
            {
                ptr=faptr;
                continue;
            }
            else //父为红，改成黑色，直接结束
            {
                faptr->RBc=0;
                break;
            }
        }
        if(outsdc&&outsdc->RBc) //兄弟有外侧红孩儿
        {
            if(islc(broptr))right_revolve(faptr,broptr);
            else left_revolve(faptr,broptr);
            outsdc->RBc=0,broptr->RBc=faptr->RBc,faptr->RBc=0;
            break;
        }
        if(insdc&&insdc->RBc) //兄弟内侧孩子为红色且外侧非红，修改成外侧红孩儿
        {
            if(islc(insdc))right_revolve(broptr,insdc);
            else left_revolve(broptr,insdc);
            broptr->RBc=1,insdc->RBc=0;
        }
    }
}

void remove_node(int val) //删除，思路是先进行值的临时交换，将删除值转移到叶节点，再进行颜色修正
{
    node* ptr=find_node(val,-1);
    node* next=NULL;
    while(ptr->lc||ptr->rc) //向下交换节点值域，而删除节点的值已无意义，故这里只要覆盖即可
    {
        if(!ptr->lc)next=ptr->rc; //只要是链型直接覆盖
        else if(!ptr->rc)next=ptr->lc;
        else next=ptr->succ();
        ptr->num--;
        ptr->val=next->val;
        ptr=next;
    }
    if(!ptr->RBc)
    {
        ptr->num--; //看似多余，但是在向上维护父节点节点域，这里这样写就方便了很多
        solve_double_black(ptr);
    }
    if(ptr==root) //特判根
    {
        root=NULL;
        delete ptr;
        return;
    }
    if(islc(ptr))ptr->fa->lc=NULL;
    else ptr->fa->rc=NULL;
    ptr->fa->maintain();
    delete ptr;
}

int main()
{
    int n;
    cin>>n;
    for(int i=0;i<n;i++)
    {
        int chose,k;
        cin>>chose>>k;
        switch(chose)
        {
            case 1:{insert_node(k);break;}
            case 2:{remove_node(k);break;}
            case 3:{cout<<get_k_rank(k,root)<<"\n";break;}
            case 4:{cout<<query_rank_k(k,root)<<"\n";break;}
            case 5:{cout<<find_lower(k)->val<<"\n";break;}
            case 6:{cout<<find_higher(k)->val<<"\n";break;}
        }
    }
    return 0;
}