PAT (Advanced Level) 1066. Root of AVL Tree (25)

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

const int maxn=30;
int n,a[maxn];
struct Node
{
    int num;
    int L,R;
    int LH,RH;
    int fa;
}node[maxn];
int sz;
int root;

void init()
{
    root=0;
    sz=0;
    node[root].fa=-1;
    node[root].L=-1,node[root].R=-1;
    node[root].num=a[1];
    node[root].LH=0,node[root].RH=0;
}

void dfs(int now)
{
    if(node[now].L!=-1) dfs(node[now].L);
    if(node[now].R!=-1) dfs(node[now].R);
    node[now].LH=max(node[node[now].L].LH,
                    node[node[now].L].RH)+1;
    node[now].RH=max(node[node[now].R].LH,
                    node[node[now].R].RH)+1;
}

void Update(int fa,int num,int tag)
{
    sz++;
    node[sz].fa=fa;
    node[sz].L=-1,node[sz].R=-1;
    node[sz].num=num;
    node[sz].LH=0,node[sz].RH=0;
    if(tag==1) node[fa].R=sz;
    else node[fa].L=sz;
    dfs(root);
}

void Insert(int num)
{
    int p=root;
    while(1)
    {
        if(num>=node[p].num)
        {
            if(node[p].R!=-1) p=node[p].R;
            else { Update(p,num,1); break; }
        }
        else
        {
            if(node[p].L!=-1) p=node[p].L;
            else { Update(p,num,0); break; }
        }
    }
}

void LL(int id)
{
    int Lson=node[id].L;

    Node tmp1=node[id];
    Node tmp2=node[Lson];

    if(tmp1.fa!=-1)
    {
        if(node[tmp1.fa].L==id) node[tmp1.fa].L=Lson;
        else node[tmp1.fa].R=Lson;
    }
    node[Lson].fa=tmp1.fa;

    node[Lson].R=id;
    node[id].fa=Lson;

    node[id].L=tmp2.R;
    node[tmp2.R].fa=id;
}

void RR(int id)
{
    int Rson=node[id].R;
    Node tmp1=node[id];
    Node tmp2=node[Rson];

    if(tmp1.fa!=-1)
    {
        if(node[tmp1.fa].L==id) node[tmp1.fa].L=Rson;
        else node[tmp1.fa].R=Rson;
    }

    node[Rson].fa=tmp1.fa;

    node[Rson].L=id;
    node[id].fa=Rson;

    node[id].R=tmp2.L;
    node[tmp2.L].fa=id;
}

void LR(int id)
{
    int Lson=node[id].L;
    RR(Lson);
    LL(id);
}

void RL(int id)
{
    int Rson=node[id].R;
    LL(Rson);
    RR(id);
}

void Rotate(int id)
{
    int need=-1;
    int p=node[id].fa;
    while(1)
    {
        if(p==-1) break;
        if(abs(node[p].LH-node[p].RH)>1) { need=p; break; }
        p=node[p].fa;
    }

    if(need==-1) return;

    if(node[need].LH>node[need].RH)
    {
        int Lson=node[need].L;
        if(node[Lson].LH>node[Lson].RH) LL(need);
        else LR(need);
    }

    else
    {
        int Rson=node[need].R;
        if(node[Rson].RH>node[Rson].LH) RR(need);
        else RL(need);
    }

    for(int i=0;i<=sz;i++) if(node[i].fa==-1) root=i;
    dfs(root);
}

int main()
{
    scanf("%d",&n);
    for(int i=1;i<=n;i++) scanf("%d",&a[i]);

    init();
    for(int i=2;i<=n;i++)
    {
        Insert(a[i]);
        Rotate(sz);
    }

    printf("%d\n",node[root].num);

    return 0;
}