1066 Root of AVL Tree

Posted 2021-03-07 feiief

link

 

 

#include <bits/stdc++.h>
# define LL long long
using namespace std;

class TreeNode{
public:
    int val;
    TreeNode* left;
    TreeNode* right;

    TreeNode(int v):val{v},left{NULL},right{NULL}{}

    int getHeight(){
        int left=this->left==NULL?0:this->left->getHeight();
        int right=this->right==NULL?0:this->right->getHeight();
        return 1+max(left,right);
    }

    int leftHeight(){
        if(this->left==NULL) return 0;
        return this->left->getHeight();
    }

    int rightHeight(){
        if(this->right==NULL) return 0;
        return this->right->getHeight();
    }

    void add(int n){
        if(n<this->val){
            if(this->left==NULL){
                this->left=new TreeNode(n);
            }else{
                this->left->add(n);
            }
        }else{
            if(this->right==NULL){
                this->right=new TreeNode(n);
            }else{
                this->right->add(n);
            }
        }
        if(this->leftHeight()-this->rightHeight()>1){
            if(this->left->rightHeight()>this->left->leftHeight()){
                this->left->leftRotate();
            }
            this->rightRotate();
            return;
        }

        if(this->rightHeight()-this->leftHeight()>1){
            if(this->right->leftHeight()>this->right->rightHeight()){
                this->right->rightRotate();
            }
            this->leftRotate();
        }
    }

    void leftRotate(){
        TreeNode* newNode=new TreeNode(this->val);
        newNode->left=this->left;
        newNode->right=this->right->left;
        this->val=this->right->val;
        this->left=newNode;
        this->right=this->right->right;
    }

    void rightRotate(){
        TreeNode* newNode=new TreeNode(this->val);
        newNode->right=this->right;
        newNode->left=this->left->right;
        this->val=this->left->val;
        this->right=newNode;
        this->left=this->left->left;
    }
};

int main(){
    int N;
    scanf("%d", &N);
    int r=0;
    scanf("%d", &r);
    TreeNode* root=new TreeNode(r);
    for(int i=1;i<=N-1;i++){
        scanf("%d", &r);
        root->add(r);
    }
    cout<<root->val;
    return 0;
}