PAT004 Root of AVL Tree

题目：

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 <stdio.h>
typedef struct treeNode {
    int data;
    struct treeNode *left;
    struct treeNode *right;
    int height;
} AVLTreeNode;

// 在PAT提交时出现MAX宏未定义的编译错误，故添加以下几行代码
#ifndef MAX
#define MAX(A, B) ((A) > (B) ? (A) : (B))
#endif

// 获取节点高度
int GetHeight(AVLTreeNode *treeNode)
{
    if (!treeNode) {
        return 0;
    } else {
        return MAX(GetHeight(treeNode->left), GetHeight(treeNode->right)) + 1;
    }
}

AVLTreeNode *SingleLeftRotation(AVLTreeNode *A)
{
    AVLTreeNode *B = A->left;
    A->left = B->right;
    B->right = A;
    A->height = MAX(GetHeight(A->left), GetHeight(A->right)) + 1;
    B->height = MAX(GetHeight(B->left), GetHeight(B->right)) + 1;
    return B;
}

AVLTreeNode *SingleRightRotation(AVLTreeNode *A)
{
    AVLTreeNode *B = A->right;
    A->right = B->left;
    B->left = A;
    A->height = MAX(GetHeight(A->left), GetHeight(A->right)) + 1;
    B->height = MAX(GetHeight(B->left), GetHeight(B->right)) + 1;
    return B;
}

AVLTreeNode *DoubleLeftRightRotation(AVLTreeNode *A)
{
    A->left = SingleRightRotation(A->left);
    return SingleLeftRotation(A);
}

AVLTreeNode *DoubleRightLeftRotation(AVLTreeNode *A)
{
    A->right = SingleLeftRotation(A->right);
    return SingleRightRotation(A);
}

// 将data插入到AVL树tree中，并返回调整后的AVL树
AVLTreeNode *AVL_insertion(int data, AVLTreeNode *tree)
{
    if (!tree) { // 若插入到空树中，新建一个节点
        tree = (AVLTreeNode *)malloc(sizeof(AVLTreeNode));
        tree->data = data;
        tree->height = 0;
        tree->left = tree->right = NULL;
    } else if (data < tree->data) { // 插入到左子树中
        tree->left = AVL_insertion(data, tree->left);
        if (GetHeight(tree->left) - GetHeight(tree->right) == 2) { // 需要左旋
            if (data < tree->left->data) { // 左单旋
                tree = SingleLeftRotation(tree);
            } else { // 左右双旋
                tree = DoubleLeftRightRotation(tree);
            }
        }
    } else if (data > tree->data) { // 插入到右子树中
        tree->right = AVL_insertion(data, tree->right);
        if (GetHeight(tree->right) - GetHeight(tree->left) == 2) { // 需要右旋
            if (data > tree->right->data) { //右单旋
                tree = SingleRightRotation(tree);
            } else {
                tree = DoubleRightLeftRotation(tree); // 右左旋
            }
        }
    } /* else data == tree->data 无需插入*/
    
    tree->height = MAX(GetHeight(tree->left), GetHeight(tree->right)) + 1;
    
    return tree;
}

int main()
{
    // 读取输入
    int count = 0;
    scanf("%d", &count);
    
    AVLTreeNode *tree = NULL;
    for (int i = 0; i < count; i++) {
        int data = 0;
        scanf("%d", &data);
        tree = AVL_insertion(data, tree);
    }
    printf("%d", tree->data);
}

运行结果：