Skip to content

Instantly share code, notes, and snippets.

@jacky860226
Created April 25, 2016 08:20
Show Gist options
  • Save jacky860226/23b096f7075b490957d75eaca3625522 to your computer and use it in GitHub Desktop.
Save jacky860226/23b096f7075b490957d75eaca3625522 to your computer and use it in GitHub Desktop.
Red Black Tree
/* 0 => black , 1 => red */
#ifndef SYMMETRIC_BINARY_B_TREE
#define SYMMETRIC_BINARY_B_TREE
template<typename T>
class rb_tree{
private:
struct node{
node *ch[2],*pa;
int s;
bool c;
T data;
node(const T&d,node*p):pa(p),s(1),c(1),data(d){}
node():s(0),c(0){ch[0]=ch[1]=pa=this;}
}*nil,*root;
inline void rotate(node*o,bool d){
if(!o->pa->s)root=o->ch[!d];
else o->pa->ch[o==o->pa->ch[1]]=o->ch[!d];
o->ch[!d]->pa=o->pa;
o->pa=o->ch[!d];
o->ch[!d]=o->pa->ch[d];
o->pa->ch[d]=o;
if(o->ch[!d]->s)o->ch[!d]->pa=o;
o->pa->s=o->s;
o->s=o->ch[0]->s+o->ch[1]->s+1;
}
inline node *find(node*o,const T&data){
while(o->s&&o->data!=data)o=o->ch[o->data<data];
return o->s?o:0;
}
inline void insert_fix(node*&o){
while(o->pa->c){
bool d=o->pa==o->pa->pa->ch[0];
node *uncle=o->pa->pa->ch[d];
if(uncle->c){
uncle->c=o->pa->c=0;
o->pa->pa->c=1;
o=o->pa->pa;
}else{
if(o==o->pa->ch[d]){
o=o->pa;
rotate(o,!d);
}
o->pa->c=0;
o->pa->pa->c=1;
rotate(o->pa->pa,d);
}
}
root->c=0;
}
inline void erase_fix(node*&x){
while(x!=root&&!x->c){
bool d=x==x->pa->ch[0];
node *w=x->pa->ch[d];
if(w->c){
w->c=0;
x->pa->c=1;
rotate(x->pa,!d);
w=x->pa->ch[d];
}
if(!w->ch[0]->c&&!w->ch[1]->c){
w->c=1,x=x->pa;
}else{
if(!w->ch[d]->c){
w->ch[!d]->c=0;
w->c=1;
rotate(w,d);
w=x->pa->ch[d];
}
w->c=x->pa->c;
w->ch[d]->c=x->pa->c=0;
rotate(x->pa,!d);
break;
}
}
x->c=0;
}
void clear(node*&o){
if(o->s)clear(o->ch[0]),clear(o->ch[1]),delete(o);
}
public:
rb_tree():nil(new node),root(nil){}
~rb_tree(){clear(root),delete nil;}
inline void clear(){clear(root),root=nil;}
inline void insert(const T&data){
node *o=root;
if(!o->s){
root=new node(data,nil);
root->ch[0]=root->ch[1]=nil;
}else{
for(;;){
o->s++;
bool d=o->data<data;
if(o->ch[d]->s)o=o->ch[d];
else{
o->ch[d]=new node(data,o),o=o->ch[d];
o->ch[0]=o->ch[1]=nil;
break;
}
}
}
insert_fix(o);
}
inline bool erase(const T&data){
node *o=find(root,data);
if(!o)return 0;
node *t=o,*p;
if(o->ch[0]->s&&o->ch[1]->s){
t=o->ch[1];
while(t->ch[0]->s)t=t->ch[0];
}
p=t->ch[!t->ch[0]->s];
p->pa=t->pa;
if(!t->pa->s)root=p;
else t->pa->ch[t->pa->ch[1]==t]=p;
if(t!=o)o->data=t->data;
for(o=t->pa;o->s;o=o->pa)o->s--;
if(!t->c)erase_fix(p);
delete t;
return 1;
}
inline bool find(const T&data){
return find(root,data);
}
inline int rank(const T&data){
int cnt=0;
for(node *o=root;o->s;)
if(o->data<data)cnt+=o->ch[0]->s+1,o=o->ch[1];
else o=o->ch[0];
return cnt;
}
inline const T&kth(int k){
for(node *o=root;;)
if(k<=o->ch[0]->s)o=o->ch[0];
else if(k==o->ch[0]->s+1)return o->data;
else k-=o->ch[0]->s+1,o=o->ch[1];
}
inline const T&operator[](int k){
return kth(k);
}
inline const T&preorder(const T&data){
node *x=root,*y=0;
while(x->s)
if(x->data<data)y=x,x=x->ch[1];
else x=x->ch[0];
if(y)return y->data;
return data;
}
inline const T&successor(const T&data){
node *x=root,*y=0;
while(x->s)
if(data<x->data)y=x,x=x->ch[0];
else x=x->ch[1];
if(y)return y->data;
return data;
}
inline int size(){return root->s;}
};
#endif
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment