PTA甲级 1066 Root of AVL Tree (C++)

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(\le 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 the 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

Solution:

// Talk is cheap, show me the code
// Created by Misdirection 2021-08-20 11:58:29
// All rights reserved.

#include <iostream>
#include <vector>

using namespace std;

struct Node{
    int key;
    Node* left;
    Node* right;
    int height;

    Node(int k = -1, Node* l = NULL, Node* r = NULL){
        key = k;
        left = l;
        right = r;
        height = 1;
    }

    ~Node(){}
};

Node* LL(Node* r){
    Node *tmp = r -> left;
    r -> left = tmp -> right;
    tmp -> right = r;

    return tmp;
}

Node* RR(Node *r){
    Node *tmp = r -> right;
    r -> right = tmp -> left;
    tmp -> left = r;

    return tmp;
}

Node* LR(Node* r){
    r -> left = RR(r -> left);
    return LL(r);
}

Node* RL(Node* r){
    r -> right = LL(r -> right);
    return RR(r);
}

int height(Node *r){
    if(r == NULL) return 0;
    else return max(height(r -> left), height(r -> right)) + 1;
}

Node* insert(Node* r, int key){
    // 分成几种情况：
    // 左孩子的左子树——LL
    // 右孩子的右子树——RR
    // 右孩子的左子树——RL
    // 左孩子的右子树——LR

    if(r == NULL){
        r = new Node(key, NULL, NULL);
        return r;
    }

    if(key < r -> key){
        r -> left = insert(r -> left, key);

        // 插入了左孩子
        if(height(r -> left) - height(r -> right) == 2){
            if(key < r -> left -> key) r = LL(r);
            else r = LR(r);

            return r;
        }
    }
    else if(key > r -> key){
        r -> right = insert(r -> right, key);

        // 插入了右孩子
        if(height(r -> right) - height(r -> left) == 2){
            if(key > r -> right -> key) r = RR(r);
            else r = RL(r);

            return r;
        }
    }

    return r;
}



int main(){
    
    int n;

    scanf("%d", &n);
    Node *root;

    for(int i = 0; i < n; ++i){
        int tmp;
        scanf("%d", &tmp);
        
        if(i == 0) root = new Node(tmp, NULL, NULL);

        else root = insert(root, tmp);
    }

    printf("%d", root -> key);

    return 0;
}