1066. Root of AVL Tree (25)

1066. Root of AVL Tree (25)

时间限制

100 ms

内存限制

65536 kB

代码长度限制

16000 B

判题程序

Standard

作者

CHEN, Yue

An AVL tree is a self-balancing binary search tree. In an AVL tree, the heights of the two child subtrees of any node differ by at most one; if at any time they differ by more than one, rebalancing is done to restore this property. Figures 1-4 illustrate the rotation rules.

    

    

Now given a sequence of insertions, you are supposed to tell the root of the resulting AVL tree.

Input Specification:

Each input file contains one test case. For each case, the first line contains a positive integer N (<=20) which is the total number of keys to be inserted. Then N distinct integer keys are given in the next line. All the numbers in a line are separated by a space.

Output Specification:

For each test case, print ythe root of the resulting AVL tree in one line.

Sample Input 1:

5
88 70 61 96 120

Sample Output 1:

70

Sample Input 2:

7
88 70 61 96 120 90 65

Sample Output 2:

88

#include <cstdio>
#include <iostream>
#include <cassert>

struct Node{
    int value;
    int height;
    Node* left;
    Node* right;
    Node(int v):left(NULL), right(NULL), height(0), value(v){}
    Node():left(NULL), right(NULL), height(0), value(0){}
};

Node* insert(Node* root, Node* newNode);

int main(){
    int count, value;
    Node* root = NULL;
    std::cin >> count;
    while(count--){
        std::cin >> value;
        root = insert(root, new Node(value));
    }
    printf("%d\n", root -> value);
    return 0;
}

int getHeight(Node* node){
    if(node == NULL)
        return -1;
    else
        return node -> height;
}

int max(int a, int b){
    if(a > b){
        return a;
    }
    else{
        return b;
    }
}

Node* singleRotateLeft(Node* r){
    Node* newRoot = r -> left;
    r -> left = newRoot -> right;
    newRoot -> right = r;
    r -> height = max(getHeight(r -> left), getHeight(r -> right)) + 1;
    newRoot -> height = max(getHeight(newRoot -> left), getHeight(newRoot -> right)) + 1;
    return newRoot;
}

Node* doubleRotateLeft(Node* r){
    Node* newRoot = r -> left -> right;
    r -> left -> right = newRoot -> left;
    newRoot -> left = r -> left;
    r -> left = newRoot -> right;
    newRoot -> right = r;
    newRoot -> left -> height = max(getHeight(newRoot -> left -> left), getHeight(newRoot -> left -> right)) + 1;
    newRoot -> right -> height = max(getHeight(newRoot -> right -> left), getHeight(newRoot -> right -> right)) + 1;
    newRoot -> height = max(getHeight(newRoot -> left), getHeight(newRoot -> right)) + 1;
    return newRoot;
}

Node* singleRotateRight(Node* r){
    Node* newRoot = r -> right;
    r -> right = newRoot -> left;
    newRoot -> left = r;
    r -> height = max(getHeight(r -> left), getHeight(r -> right)) + 1;
    newRoot -> height = max(getHeight(newRoot -> left), getHeight(newRoot -> right)) + 1;
    return newRoot;
}

Node* doubleRotateRight(Node* r){
    Node* newRoot = r -> right -> left;
    r -> right -> left = newRoot -> right;
    newRoot -> right = r -> right;
    r -> right = newRoot -> left;
    newRoot -> left = r;
    newRoot -> left -> height = max(getHeight(newRoot -> left -> left), getHeight(newRoot -> left -> right)) + 1;
    newRoot -> right -> height = max(getHeight(newRoot -> right -> left), getHeight(newRoot -> right -> right)) + 1;
    newRoot -> height = max(getHeight(newRoot -> left), getHeight(newRoot -> right)) + 1;
    return newRoot;
}

Node* insert(Node* root, Node* newNode){
    Node* r = root;
    if (r == NULL){
        r = newNode;
    }
    else{
        if (newNode -> value < r -> value){
            r -> left = insert(r -> left, newNode);
            if(getHeight(r -> left) - getHeight(r -> right) >= 2){
                if(newNode -> value < r -> left -> value)
                    r = singleRotateLeft(r);
                else
                    r = doubleRotateLeft(r);
            }
        }
        else if(newNode -> value > r -> value){
            r -> right = insert(r -> right, newNode);
            if(getHeight(r -> right) - getHeight(r -> left) >= 2){
                if(newNode -> value > r -> right -> value)
                    r = singleRotateRight(r);
                else
                    r = doubleRotateRight(r);
            }
        }
        else{
            assert(0);
        }
        r -> height = max(getHeight(r -> left), getHeight(r -> right)) + 1;
    }
    return r;
}