SBT.cpp

// Size-Balanced Tree, taken from http://blog.tomtung.com/2007/06/size-balanced-tree-in-cpp
// TODO: How does this behave with duplicate keys?
struct SBTNode{
    SBTNode *ch[2],*p;
    long key;
    unsigned long size;
    SBTNode(long _key,unsigned long _size);
}NIL=SBTNode(0,0);
// To make an empty tree: SBTree t = &NIL;
typedef SBTNode *SBTree;
SBTNode::SBTNode(long _key,unsigned long _size=1){
    ch[0]=ch[1]=p=&NIL;
    size=_size;
    key=_key;
}

// Returns node if key is found, &NIL otherwise
SBTNode *SBT_Search(SBTree T,long key);

// Example invocation: SBT_Insert(t, new SBTNode(42));
void SBT_Insert(SBTree &T, SBTNode* x);

// Careful, if no matching key is found, this will still delete a node from the tree.
// The deleted node is not freed, and a pointer to it is returned.
SBTNode *SBT_Delete(SBTree &T, long key);

// Returns first node whose key is strictly less than key, or &NIL if no such node
SBTNode *SBT_Pred(SBTree T, long key);

// Returns first node whose key is strictly greater than key, or &NIL if no such node
SBTNode *SBT_Succ(SBTree T,long key);

// Returns node with ith smallest value (1 <= i <= T->size)
SBTNode *SBT_Select(SBTree T, unsigned long i);

// Returns k such that key is the kth smallest key in the tree (1 <= k <= T->size), or 0 if key not found.
unsigned long SBT_Rank(SBTree T, long key);

inline void SBT_Rotate(SBTree &x,bool flag){
    SBTNode *y=x->ch[!flag];
	// assert(x!=&NIL&&y!=&NIL);
    y->p=x->p;
    x->p=y;
    if(y->ch[flag]!=&NIL) y->ch[flag]->p=x;
    x->ch[!flag]=y->ch[flag];
    y->ch[flag]=x;
    y->size=x->size;
    x->size=x->ch[0]->size+x->ch[1]->size+1;
    x=y;
}
void SBT_Maintain(SBTree &T,bool flag){
    if(T->ch[flag]->ch[flag]->size>T->ch[!flag]->size)
        SBT_Rotate(T,!flag);
    else if(T->ch[flag]->ch[!flag]->size>T->ch[!flag]->size){
        SBT_Rotate(T->ch[flag],flag);
        SBT_Rotate(T,!flag);
    }
    else return;
    SBT_Maintain(T->ch[0],0),SBT_Maintain(T->ch[1],1);
    SBT_Maintain(T,0),SBT_Maintain(T,1);
}
 
SBTNode *SBT_Search(SBTree T,long key){
    return T==&NIL||T->key==key?T:SBT_Search(T->ch[key>T->key],key);
}

void SBT_Insert(SBTree &T, SBTNode* x){
    if(T==&NIL)    T=x;
    else{
        T->size++;
        x->p=T;
        SBT_Insert(T->ch[x->key>T->key],x);
        SBT_Maintain(T,x->key>T->key);
    }
}
 
SBTNode *SBT_Delete(SBTree &T, long key){
    if(T==&NIL)    return &NIL;
    T->size--;
    if(T->key==key||T->ch[key>T->key]==&NIL){
        SBTNode *toDel;
        if(T->ch[0]==&NIL||T->ch[1]==&NIL){
            toDel=T;
            T=T->ch[T->ch[1]!=&NIL];
            if(T!=&NIL)    T->p=toDel->p;
        }else{
            toDel=SBT_Delete(T->ch[1],key-1);
            T->key=toDel->key;
        }
        return toDel;
    }
    else return SBT_Delete(T->ch[key>T->key],key);
}
 
SBTNode *SBT_Pred(SBTree T, long key){
    if(T==&NIL)    return &NIL;
    if(key<=T->key)    return SBT_Pred(T->ch[0],key);
    else{
        SBTNode *pred=SBT_Pred(T->ch[1],key);
        return (pred!=&NIL?pred:T);
    }
}
 
SBTNode *SBT_Succ(SBTree T,long key){
    if(T==&NIL)    return &NIL;
    if(key>=T->key)    return SBT_Succ(T->ch[1],key);
    else{
        SBTNode *succ= SBT_Succ(T->ch[0],key);
        return(succ!=&NIL?succ:T);
    }
}
 
SBTNode *SBT_Select(SBTree T, unsigned long i){
    unsigned long r = T->ch[0]->size+1;
    if(i==r)    return T;
    else return SBT_Select(T->ch[i>r],i>r?i-r:i);
}
 
unsigned long SBT_Rank(SBTree T, long key){
    if(T==&NIL)    return 0;
    if(T->key==key)    return T->ch[0]->size+1;
    else if(key<T->key)    return SBT_Rank(T->ch[0],key);
    else{
        unsigned long r=SBT_Rank(T->ch[1],key);
        return r==0?0:r+T->ch[0]->size+1;
    }
}