Arne Andersson Tree

aa_tree.cpp

#ifndef ARNE_ANDERSSON_TREE
#define ARNE_ANDERSSON_TREE
template<typename T>
class aa_tree{
	private:
		struct node{
			node *ch[2];
			int s,level;
			T data;
			node(const T&d):s(1),level(1),data(d){}
			node():s(0),level(0){ch[0]=ch[1]=this;}
		}*nil,*root;
		inline void rotate(node *&a,bool d){
			if(a->level==a->ch[d]->ch[d]->level){
				node *b=a;
				a=a->ch[d];
				b->ch[d]=a->ch[!d];
				a->ch[!d]=b;
				if(d)a->level++;
				a->s=b->s;
				b->s=b->ch[0]->s+b->ch[1]->s+1;
			}
		}
		inline void erase_fix(node *&o){
			if(o ->ch[0]->level<o->level-1||o->ch[1]->level<o->level-1){
				if(o->ch[1]->level>--o->level)o->ch[1]->level=o->level;
				rotate(o,0);
				if(o->ch[1]->s)rotate(o->ch[1],0);
				if(o->ch[1]->ch[1]->s)rotate(o->ch[1]->ch[1],0);
				rotate(o,1);
				if(o->ch[1]->s)rotate(o->ch[1],1);
			}
		}
		void insert(node *&o,const T&data){
			if(!o->s){
				o=new node(data);
				o->ch[0]=o->ch[1]=nil;
			}else{
				o->s++;
				insert(o->ch[o->data<data],data);
				rotate(o,0);
				rotate(o,1); 
			}
		}
		bool erase(node *&o,const T&data){
			if(!o->s)return 0;
			if(o->data==data){
				if(!o->ch[0]->s||!o->ch[1]->s){
					node *t=o;
					o=o->ch[0]->s?o->ch[0]:o->ch[1];
					delete t;
				}else{
					o->s--;
					node *t=o->ch[1];
					while(t->ch[0]->s)t=t->ch[0];
					o->data=t->data;
					erase(o->ch[1],t->data);
					erase_fix(o);
				}return 1;
			}else if(erase(o->ch[o->data<data],data)){
				o->s--,erase_fix(o);
				return 1;
			}else return 0;
		}
		void clear(node *&o){
			if(o->s)clear(o->ch[0]),clear(o->ch[1]),delete(o);
		}
	public:
		aa_tree():nil(new node){
			root=nil->ch[0]=nil->ch[1]=nil;
		}
		~aa_tree(){clear(root),delete nil;}
		inline void clear(){clear(root),root=nil;}
		inline void insert(const T&data){
			insert(root,data);
		}
		inline bool erase(const T&data){
			return erase(root,data);
		}
		inline bool find(const T&data){
			node *o=root;
			while(o->s&&o->data!=data)o=o->ch[o->data<data];
			return o->s?1:0;
		}
		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