【PAT】1123 Is It a Complete AVL Tree (30 分)

#include <iostream>
#include <math.h>
#include <queue>
using namespace std;

struct node
{
    int data;
    node* lchild,* rchild;
};

int getHeight(node* root){
    if(!root) return 0;
    int l=getHeight(root->lchild);
    int r=getHeight(root->rchild);
    return max(l,r)+1;
}

node* rightRotate(node* root){
    node* B=root->lchild;
    root->lchild=B->rchild;
    B->rchild=root;
    //如果每次新调用一次getHeight()函数,则不必更新树高~
    return B;
}
node* leftRotate(node* root){
    node* B=root->rchild;
    root->rchild=B->lchild;
    B->lchild=root;
    //如果每次新调用一次getHeight()函数,则不必更新树高~
    return B;
}

node* leftRightRotate(node* root){
    root->lchild=leftRotate(root->lchild);//把末端拆了,换个位置,再重接上去
    return rightRotate(root);//再换一次
}

node* rightLeftRotate(node* root){
    root->rchild=rightRotate(root->rchild);
    return leftRotate(root);
}

node* InsertT(node* root,int v){
    if(!root){
        root=new node;
        root->data=v;
        root->lchild=root->rchild=NULL;
        return root;
    }

    if(v<root->data){//左子树
        root->lchild=InsertT(root->lchild,v);
        if(getHeight(root->lchild)-getHeight(root->rchild)>=2){//左子树异常->右旋
            if(v<root->lchild->data){//左子树的左子树->RR
                root=rightRotate(root);
            }else{//左子树的右子树->LR
                root=leftRightRotate(root);
            }
        }
    }else{//右子树
        root->rchild=InsertT(root->rchild,v);
        if(getHeight(root->rchild)-getHeight(root->lchild)>=2){//右子树异常->左旋
            if(v>root->rchild->data){//右子树的右子树->LL
                root=leftRotate(root);
            }else{//右子树的左子树->RL
                root=rightLeftRotate(root);
            }
        }
    }
    return root;  
}

int flag=0;
int after=1,isComplete=1;
/*不是完全二叉树只有一种情况:左子树不存在,右子树存在
完全二叉树只可能用到结尾左子树存在,右子树不存在的一种判定*/
void levelTraversal(node* root){
    if(!root) return;
    queue<node*> q;
    q.push(root);
    while(!q.empty()){
        node* tmp=q.front();
        q.pop();
        printf("%s%d",flag==1?" ":"",tmp->data);
        flag=1;
        if(tmp->lchild){
            if(after==0) isComplete=0;
            q.push(tmp->lchild);
        }else{
            after=0;
        }
        if(tmp->rchild){
            if(after==0) isComplete=0;
            q.push(tmp->rchild);
        }else{
            after=0;
        }
    }
}

int main(){
    int n;
    scanf("%d",&n);
    int a[n];
    node* root=NULL;
    for(int i=0;i<n;i++){
        scanf("%d",&a[i]);
        root=InsertT(root,a[i]);
    }
    levelTraversal(root);
    printf("\n%s",isComplete?"YES":"NO");

    return 0;
}