AVL平衡树

#include <iostream>
#include <algorithm>

// AVL树节点结构
struct AVLNode {
    int key;
    AVLNode* left;
    AVLNode* right;
    int height;

    AVLNode(int k) : key(k), left(nullptr), right(nullptr), height(1) {}
};

class AVLTree {
private:
    AVLNode* root;

    // 获取节点高度
    int height(AVLNode* node) {
        return node ? node->height : 0;
    }

    // 计算平衡因子
    int getBalance(AVLNode* node) {
        return node ? height(node->left) - height(node->right) : 0;
    }

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

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

        // 更新高度
        y->height = std::max(height(y->left), height(y->right)) + 1;
        x->height = std::max(height(x->left), height(x->right)) + 1;

        return x;
    }

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

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

        // 更新高度
        x->height = std::max(height(x->left), height(x->right)) + 1;
        y->height = std::max(height(y->left), height(y->right)) + 1;

        return y;
    }

    // 插入节点辅助函数
    AVLNode* insert(AVLNode* node, int key) {
        // 1. 执行标准BST插入
        if (!node) return new AVLNode(key);

        if (key < node->key)
            node->left = insert(node->left, key);
        else if (key > node->key)
            node->right = insert(node->right, key);
        else // 不允许重复键
            return node;

        // 2. 更新节点高度
        node->height = 1 + std::max(height(node->left), height(node->right));

        // 3. 获取平衡因子检查是否平衡
        int balance = getBalance(node);

        // 如果不平衡，有4种情况

        // 左左情况
        if (balance > 1 && key < node->left->key)
            return rightRotate(node);

        // 右右情况
        if (balance < -1 && key > node->right->key)
            return leftRotate(node);

        // 左右情况
        if (balance > 1 && key > node->left->key) {
            node->left = leftRotate(node->left);
            return rightRotate(node);
        }

        // 右左情况
        if (balance < -1 && key < node->right->key) {
            node->right = rightRotate(node->right);
            return leftRotate(node);
        }

        return node;
    }

    // 查找最小节点
    AVLNode* minValueNode(AVLNode* node) {
        AVLNode* current = node;
        while (current->left)
            current = current->left;
        return current;
    }

    // 删除节点辅助函数
    AVLNode* deleteNode(AVLNode* root, int key) {
        // 1. 执行标准BST删除
        if (!root) return root;

        if (key < root->key)
            root->left = deleteNode(root->left, key);
        else if (key > root->key)
            root->right = deleteNode(root->right, key);
        else {
            // 节点有一个或没有子节点
            if (!root->left || !root->right) {
                AVLNode* temp = root->left ? root->left : root->right;

                // 没有子节点的情况
                if (!temp) {
                    temp = root;
                    root = nullptr;
                } else // 一个子节点的情况
                    *root = *temp; // 复制非空子节点内容

                delete temp;
            } else {
                // 节点有两个子节点：获取中序后继（右子树的最小值）
                AVLNode* temp = minValueNode(root->right);

                // 复制中序后继的数据到当前节点
                root->key = temp->key;

                // 删除中序后继
                root->right = deleteNode(root->right, temp->key);
            }
        }

        // 如果树只有一个节点则返回
        if (!root) return root;

        // 2. 更新当前节点的高度
        root->height = 1 + std::max(height(root->left), height(root->right));

        // 3. 获取平衡因子检查是否平衡
        int balance = getBalance(root);

        // 如果不平衡，有4种情况

        // 左左情况
        if (balance > 1 && getBalance(root->left) >= 0)
            return rightRotate(root);

        // 左右情况
        if (balance > 1 && getBalance(root->left) < 0) {
            root->left = leftRotate(root->left);
            return rightRotate(root);
        }

        // 右右情况
        if (balance < -1 && getBalance(root->right) <= 0)
            return leftRotate(root);

        // 右左情况
        if (balance < -1 && getBalance(root->right) > 0) {
            root->right = rightRotate(root->right);
            return leftRotate(root);
        }

        return root;
    }

    // 中序遍历辅助函数
    void inorder(AVLNode* node) {
        if (node) {
            inorder(node->left);
            std::cout << node->key << " ";
            inorder(node->right);
        }
    }

    // 前序遍历辅助函数
    void preorder(AVLNode* node) {
        if (node) {
            std::cout << node->key << " ";
            preorder(node->left);
            preorder(node->right);
        }
    }

public:
    AVLTree() : root(nullptr) {}

    // 插入节点
    void insert(int key) {
        root = insert(root, key);
    }

    // 删除节点
    void deleteKey(int key) {
        root = deleteNode(root, key);
    }

    // 中序遍历
    void inorder() {
        inorder(root);
        std::cout << std::endl;
    }

    // 前序遍历
    void preorder() {
        preorder(root);
        std::cout << std::endl;
    }

    // 查找节点
    AVLNode* search(int key) {
        AVLNode* current = root;
        while (current) {
            if (key < current->key)
                current = current->left;
            else if (key > current->key)
                current = current->right;
            else
                return current;
        }
        return nullptr;
    }
};

// 测试代码
int main() {
    AVLTree tree;

    // 测试插入
    tree.insert(10);
    tree.insert(20);
    tree.insert(30);
    tree.insert(40);
    tree.insert(50);
    tree.insert(25);

    std::cout << "前序遍历: ";
    tree.preorder(); // 应显示平衡后的树结构

    std::cout << "中序遍历: ";
    tree.inorder(); // 应显示有序序列

    // 测试删除
    tree.deleteKey(20);
    std::cout << "删除20后的中序遍历: ";
    tree.inorder();

    // 测试查找
    AVLNode* found = tree.search(30);
    if (found)
        std::cout << "找到节点30\n";
    else
        std::cout << "未找到节点30\n";

    return 0;
}