【PTA数据结构 | C语言版】AVL树的操作

本专栏持续输出数据结构题目集，欢迎订阅。

题目

请编写程序，实现 AVL 树的插入及旋转操作。

输入格式：
输入首先给出一个正整数 n（≤20），随后一行给出 n 个不重复的整数。数字间以空格分隔。

输出格式：
将给定的 n 个不重复的整数顺次插入一棵初始为空的 AVL 树。随后输出该树的先序遍历序列，每个数字占一行。

输入样例：
10
8 13 20 7 6 9 11 10 4 12

输出样例：
8
6
4
7
11
10
9
13
12
20

代码

#include <stdio.h>
#include <stdlib.h>

// AVL树节点结构
typedef struct Node {
    int data;
    struct Node *left;
    struct Node *right;
    int height; // 节点高度，用于平衡判断
} Node;

// 创建新节点
Node* createNode(int data) {
    Node* node = (Node*)malloc(sizeof(Node));
    node->data = data;
    node->left = NULL;
    node->right = NULL;
    node->height = 1; // 新节点初始高度为1
    return node;
}

// 获取节点高度
int getHeight(Node* node) {
    if (node == NULL)
        return 0;
    return node->height;
}

// 计算两个数中的最大值
int max(int a, int b) {
    return (a > b) ? a : b;
}

// 计算平衡因子（左子树高度 - 右子树高度）
int getBalanceFactor(Node* node) {
    if (node == NULL)
        return 0;
    return getHeight(node->left) - getHeight(node->right);
}

// 更新节点高度
void updateHeight(Node* node) {
    if (node != NULL) {
        node->height = 1 + max(getHeight(node->left), getHeight(node->right));
    }
}

// 右旋转操作
Node* rightRotate(Node* y) {
    Node* x = y->left;
    Node* T2 = x->right;

    // 执行旋转
    x->right = y;
    y->left = T2;

    // 更新高度
    updateHeight(y);
    updateHeight(x);

    // 返回新的根节点
    return x;
}

// 左旋转操作
Node* leftRotate(Node* x) {
    Node* y = x->right;
    Node* T2 = y->left;

    // 执行旋转
    y->left = x;
    x->right = T2;

    // 更新高度
    updateHeight(x);
    updateHeight(y);

    // 返回新的根节点
    return y;
}

// 插入节点到AVL树
Node* insert(Node* node, int data) {
    // 1. 执行标准的BST插入
    if (node == NULL)
        return createNode(data);

    if (data < node->data)
        node->left = insert(node->left, data);
    else if (data > node->data)
        node->right = insert(node->right, data);
    else // 不允许插入重复值
        return node;

    // 2. 更新当前节点的高度
    updateHeight(node);

    // 3. 计算平衡因子，检查是否失衡
    int balance = getBalanceFactor(node);

    // 4. 如果失衡，有四种情况需要处理

    // 左左情况（LL）- 右旋转
    if (balance > 1 && data < node->left->data)
        return rightRotate(node);

    // 右右情况（RR）- 左旋转
    if (balance < -1 && data > node->right->data)
        return leftRotate(node);

    // 左右情况（LR）- 先左旋转再右旋转
    if (balance > 1 && data > node->left->data) {
        node->left = leftRotate(node->left);
        return rightRotate(node);
    }

    // 右左情况（RL）- 先右旋转再左旋转
    if (balance < -1 && data < node->right->data) {
        node->right = rightRotate(node->right);
        return leftRotate(node);
    }

    // 如果平衡，返回当前节点
    return node;
}

// 先序遍历AVL树并打印节点值
void preOrder(Node* root) {
    if (root != NULL) {
        printf("%d\n", root->data);
        preOrder(root->left);
        preOrder(root->right);
    }
}

// 主函数
int main() {
    Node* root = NULL;
    int n, i, data;

    // 读取输入
    scanf("%d", &n);
    for (i = 0; i < n; i++) {
        scanf("%d", &data);
        root = insert(root, data);
    }

    // 输出先序遍历结果
    preOrder(root);

    return 0;
}