1066. Root of AVL Tree (25)

1066. Root of AVL Tree (25)

#include <bits/stdc++.h>
using namespace std;
struct TreeNode
{
    int val;
    TreeNode *left,*right;
    TreeNode(int v):val(v),left(NULL),right(NULL){}
};
int height(TreeNode *r)
{
    return r==NULL?0:max(height(r->left),height(r->right))+1;
}
TreeNode* lRotate(TreeNode *r)
{
    TreeNode *p=r->right;
    r->right=p->left;
    p->left=r;
    return p;
}
TreeNode* rRotate(TreeNode *r)
{
    TreeNode *p=r->left;
    r->left=p->right;
    p->right=r;
    return  p;
}
TreeNode* lrRotate(TreeNode *r)
{
    r->left=lRotate(r->left);
    return rRotate(r);
}
TreeNode* rlRotate(TreeNode *r)
{
    r->right=rRotate(r->right);
    return lRotate(r);
}
TreeNode* insert(int val,TreeNode *r)
{
    if(r==NULL) return new TreeNode(val);
    if(r->val<val){
        r->right=insert(val,r->right);
        if(abs(height(r->left)-height(r->right))>1){
            r=(r->right->val<val)?lRotate(r):rlRotate(r);
        }
    }
    else{
        r->left=insert(val,r->left);
        if(abs(height(r->left)-height(r->right))>1){
            r=(r->left->val>val)?rRotate(r):lrRotate(r);
        }
    }
    return r;
}
int main()
{
    TreeNode *r=NULL;
    int n;
    scanf("%d",&n);
    for(int i=0;i<n;++i){
        int val;
        scanf("%d",&val);
        r=insert(val,r);
    }
    printf("%d\n",r->val);
    return 0;
}