AVL树的基本操作 HDU_4006 AVL树解法

题目连接：http://acm.hdu.edu.cn/showproblem.php?pid=4006

/*
AVL树的基本操作，插入、删除、查找，平均时间复杂度为：（logN）
*/
#include 
#include 
#include 
const int NN = 1000010 ;
int N,K ;
struct Node{
    int height ;        //结点的高度
    int cnt ;           //子树中结点的个数
    int value ;         //结点的值
    struct Node *left, *right ;     //结点的左、右孩子
}q[NN] ;        //AVL树的结点
typedef Node* pnode ;
pnode T ;
int cnt ;

/*
计算以T为Root的子树中共有多少个结点
*/
int calc_cnt( pnode T ){
    if( T == NULL ) return 0 ;
    else        return T->cnt ;
}
/*
计算以T为Root的子树的高度，规定：空子树高度为-1，单结点高度为0 ;
*/
int calc_hei( pnode T ){
    if( T == NULL ) return -1 ;
    else        return T->height ;
}

/*
对T子树进行一次“一字型”的一次左旋
应用于“左左”形式
*/
pnode singleRotateLeft( pnode T ){
    pnode k = T->left ;
    T->left = k->right ;
    k->right = T ;
    T->cnt = calc_cnt( T->left ) + calc_cnt( T->right ) + 1 ;
    T->height = std::max( calc_hei( T->left ) , calc_hei( T->right ) ) + 1 ;
    k->cnt = calc_cnt( k->left ) + calc_cnt( k->right ) + 1 ;
    k->height = std::max( calc_hei( k->left ) , calc_hei( k->right ) ) + 1 ;
    return k ;
}

/*
对T子树进行一次“一字型”的一次右旋
应用于“右右”形式
*/
pnode singleRotateRight( pnode T ){
    pnode k = T->right ;
    T->right = k->left ;
    k->left = T ;
    T->cnt = calc_cnt( T->left ) + calc_cnt( T->right ) + 1 ;
    T->height = std::max( calc_hei( T->left ) , calc_hei( T->right ) ) + 1 ;
    k->cnt = calc_cnt( k->left ) + calc_cnt( k->right ) + 1 ;
    k->height = std::max( calc_hei( k->left ) , calc_hei( k->right ) ) + 1 ;
    return k ;
}
/*
对T子树进行一次“之字型”的一次左双旋
适用于“右左”形式
*/
pnode doubleRotateLeft( pnode T ){
    T->right = singleRotateLeft( T->right ) ;
    return singleRotateRight( T ) ;
}
/*
对T子树进行一次“之字型”的一次右双旋
适用于“左右”形式
*/
pnode doubleRotateRight( pnode T ){
    T->left = singleRotateRight( T->left ) ;
    return singleRotateLeft( T ) ;
}

/*在T子树中插入一个关键字是v的结点*/
pnode insert( pnode T , int v ){
    if( T==NULL ){
        T = &q[cnt++] ;
        T->value = v ;
        T->cnt = 1 ;
        T->height = 0 ;
        T->left = T->right = NULL ;
        return T ;
    }
    if( v <= T->value ){
        T->left = insert( T->left , v ) ;
        if( calc_hei(T->left) - calc_hei( T->right ) == 2 ){
            if( v <= T->left->value ){
                T = singleRotateLeft( T ) ;
            }
            else
                T = doubleRotateRight( T ) ;
        }
    }
    else{
        T->right = insert( T->right , v ) ;
        if( calc_hei( T->right ) - calc_hei(T->left) == 2 ){
            if( v <= T->right->value ){
                T = doubleRotateLeft( T ) ;
            }
            else
                T = singleRotateRight( T ) ;
        }
    }
    T->cnt = calc_cnt( T->left ) + calc_cnt( T->right ) + 1 ;
    T->height = std::max( calc_hei(T->left) , calc_hei(T->right) ) + 1 ;
    return T ;
}

pnode del( pnode T , int v ){
    if( T==NULL )   return NULL ;
    if( v < T->value ){
        T->left = del( T->left , v ) ;
        if( calc_hei( T->right ) - calc_hei( T->left) == 2 ){
            if( T->right->left!=NULL && ( calc_hei(T->right->left) > calc_hei(T->right->right) ) ){
                T = doubleRotateLeft( T ) ;
            }
            else
                T = singleRotateRight( T ) ;
        }
    }
    else if( v > T->value ){
        T->right = del( T->right , v ) ;
        if( calc_hei(T->left) - calc_hei(T->right) == 2 ){
            if( T->left->right!=NULL && ( calc_hei(T->left->right) > calc_hei(T->left->left) ) ){
                T = doubleRotateRight( T ) ;
            }
            else
                T = singleRotateLeft( T ) ;
        }
    }
    else{
        if( T->left!=NULL && T->right!=NULL ){
            pnode tmp = T->right ;
            while( tmp->left != NULL )  tmp = tmp->left ;
            T->value = tmp->value ;
            T->cnt = tmp->cnt ;
            T->right = del( T->right ,tmp->value ) ;
            if( calc_hei(T->left) - calc_hei(T->right) == 2 ){
                if( T->left->right!=NULL && ( calc_hei(T->left->right) > calc_hei(T->left->left) ) ){
                    T = doubleRotateRight( T ) ;
                }
                else
                    T = singleRotateLeft( T ) ;
            }
        }
        else{
            if( T->right == NULL ){
                T = T->left ;
            }
            else
                T = T->right ;
        }
    }
    if( T == NULL ) return NULL ;
    T->cnt = calc_cnt( T->left ) + calc_cnt( T->right ) + 1 ;
    T->height = std::max( calc_hei(T->left) , calc_hei(T->right) ) + 1 ;
    return T ;
}

int query( pnode T , int n ){
    if( calc_cnt( T->left ) < n - 1 )   return query( T->right , n - calc_cnt(T->left) - 1) ;
    else if( calc_cnt( T->left ) == n-1 )   return T->value ;
    else        return query( T->left , n ) ;
}

int main(){
    char ch[5] ;
    int a , tot ;
    while( scanf("%d %d",&N,&K) == 2 ){
        cnt = 0 ; T = NULL ; tot = 0 ;
        for(int i=0;i