PAT 1123. Is It a Complete AVL Tree (30)

AVL树的插入，旋转。

#include<map>
#include<set>
#include<ctime>
#include<cmath>
#include<queue>
#include<string>
#include<stack>
#include<vector>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<functional>
using namespace std;

int n;

struct X
{
    int val,L,R,hL,hR,fa;
}s[50];
int sz,F,root=0;

void Insert(int x)
{
    int p=root;
    while(1)
    {
        if(x<s[p].val)
        {
            if(s[p].L==-1)
            {
                sz++;
                s[p].L=sz;
                s[sz].val=x;
                s[sz].fa=p;
                break;
            }
            else p=s[p].L;
        }
        else
        {
            if(s[p].R==-1)
            {
                sz++;
                s[p].R=sz;
                s[sz].val=x;
                s[sz].fa=p;
                break;
            }
            else p=s[p].R;
        }
    }
}

void dep(int x)
{
    if(s[x].L!=-1) dep(s[x].L);
    if(s[x].R!=-1) dep(s[x].R);

    if(s[x].L!=-1) s[x].hL = max(s[s[x].L].hL,s[s[x].L].hR)+1;
    else s[x].hL=0;

    if(s[x].R!=-1) s[x].hR = max(s[s[x].R].hL,s[s[x].R].hR)+1;
    else s[x].hR=0;
}

void dfs(int x)
{
    if(abs(s[x].hL-s[x].hR)>1) F=x;

    if(s[x].L!=-1) dfs(s[x].L);
    if(s[x].R!=-1) dfs(s[x].R);
}

void Left(int x)
{
    int son = s[x].R;

    X tmpA = s[x];
    X tmpB = s[son];

    s[x].R = tmpB.L;
    s[x].fa = son;

    s[son].L = x;
    s[son].fa = tmpA.fa;

    s[tmpB.L].fa = x;

    if(s[son].fa!=-1)
    {
        int Fa = s[son].fa;
        if(s[Fa].L==x) s[Fa].L=son;
        else s[Fa].R=son;
    }
}

void Right(int x)
{
    int son = s[x].L;

    X tmpA = s[x];
    X tmpB = s[son];

    s[x].L = tmpB.R;
    s[x].fa = son;

    s[son].R = x;
    s[son].fa = tmpA.fa;

    s[tmpB.R].fa = x;

    if(s[son].fa!=-1)
    {
        int Fa = s[son].fa;
        if(s[Fa].L==x) s[Fa].L=son;
        else s[Fa].R=son;
    }
}

void bfs()
{
    queue<int>Q; Q.push(root);
    vector<int>ans;

    while(!Q.empty())
    {
        int h = Q.front(); Q.pop();
        ans.push_back(s[h].val);
        if(s[h].L!=-1) Q.push(s[h].L);
        if(s[h].R!=-1) Q.push(s[h].R);
    }

    for(int i=0;i<ans.size();i++)
    {
        printf("%d",ans[i]);
        if(i<ans.size()-1) printf(" ");
        else printf("\n");
    }

}

bool fail;

void check(int x,int id)
{
    if(id>n) fail=1;
    if(s[x].L!=-1) check(s[x].L,2*id);
    if(s[x].R!=-1) check(s[x].R,2*id+1);
}

int main()
{
    scanf("%d",&n);

    for(int i=0;i<=40;i++)
    {
        s[i].val=s[i].L=s[i].R=-1;
        s[i].hL = s[i].hR = 0;
    }

    int x; scanf("%d",&x); s[0].val=x; s[0].fa=-1;

    for(int i=2;i<=n;i++)
    {
        int x; scanf("%d",&x); Insert(x);

        dep(root); F=-1; dfs(root);
        if(F==-1) continue;

        if(s[F].hL>s[F].hR&&s[s[F].L].hL>s[s[F].L].hR) Right(F);
        else if(s[F].hL<s[F].hR&&s[s[F].R].hL<s[s[F].R].hR) Left(F);
        else if(s[F].hL>s[F].hR&&s[s[F].L].hL<s[s[F].L].hR) Left(s[F].L), Right(F);
        else if(s[F].hL<s[F].hR&&s[s[F].R].hL>s[s[F].R].hR) Right(s[F].R), Left(F);

        for(int j=0;j<=sz;j++) if(s[j].fa==-1) root=j;
    }

    bfs();

    fail=0; check(root,1);
    if(fail) printf("NO\n");
    else printf("YES\n");

    return 0;
}