AVL树

admin

[C++] 纯文本查看 复制代码

#include <iostream>
using namespace std;

// AVL树节点结构
struct Node {
    int data;
    Node* left;
    Node* right;
    int height;
    
    Node(int val) : data(val), left(nullptr), right(nullptr), height(1) {}
};

// AVL树类
class AVLTree {
private:
    Node* root;
    
    // 获取节点高度
    int height(Node* node) {
        return (node == nullptr) ? 0 : node->height;
    }
    
    // 计算平衡因子
    int getBalance(Node* node) {
        return (node == nullptr) ? 0 : height(node->left) - height(node->right);
    }
    
    // 更新节点高度
    void updateHeight(Node* node) {
        if (node != nullptr) {
            node->height = 1 + max(height(node->left), height(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;
    }
    
    // 插入节点的辅助函数
    Node* insert(Node* node, int val) {
        // 1. 执行标准BST插入
        if (node == nullptr) {
            return new Node(val);
        }
        
        if (val < node->data) {
            node->left = insert(node->left, val);
        } else if (val > node->data) {
            node->right = insert(node->right, val);
        } else {
            // 不允许插入重复值
            return node;
        }
        
        // 2. 更新当前节点高度
        updateHeight(node);
        
        // 3. 获取平衡因子，检查是否失衡
        int balance = getBalance(node);
        
        // 4. 如果失衡，有四种情况需要处理
        
        // 左左情况
        if (balance > 1 && val < node->left->data) {
            return rightRotate(node);
        }
        
        // 右右情况
        if (balance < -1 && val > node->right->data) {
            return leftRotate(node);
        }
        
        // 左右情况
        if (balance > 1 && val > node->left->data) {
            node->left = leftRotate(node->left);
            return rightRotate(node);
        }
        
        // 右左情况
        if (balance < -1 && val < node->right->data) {
            node->right = rightRotate(node->right);
            return leftRotate(node);
        }
        
        // 如果未失衡，返回当前节点
        return node;
    }
    
    // 找到最小值节点
    Node* minValueNode(Node* node) {
        Node* current = node;
        while (current->left != nullptr) {
            current = current->left;
        }
        return current;
    }
    
    // 删除节点的辅助函数
    Node* deleteNode(Node* node, int val) {
        // 1. 执行标准BST删除
        if (node == nullptr) {
            return node;
        }
        
        if (val < node->data) {
            node->left = deleteNode(node->left, val);
        } else if (val > node->data) {
            node->right = deleteNode(node->right, val);
        } else {
            // 找到要删除的节点
            
            // 情况1：叶子节点或只有一个子节点
            if (node->left == nullptr || node->right == nullptr) {
                Node* temp = (node->left != nullptr) ? node->left : node->right;
                
                // 子节点为空的情况
                if (temp == nullptr) {
                    temp = node;
                    node = nullptr;
                } else { // 一个子节点的情况
                    *node = *temp; // 复制子节点内容
                }
                delete temp;
            } else {
                // 情况2：有两个子节点
                Node* temp = minValueNode(node->right);
                node->data = temp->data;
                node->right = deleteNode(node->right, temp->data);
            }
        }
        
        // 如果树只有一个节点，返回
        if (node == nullptr) {
            return node;
        }
        
        // 2. 更新当前节点高度
        updateHeight(node);
        
        // 3. 获取平衡因子，检查是否失衡
        int balance = getBalance(node);
        
        // 4. 如果失衡，有四种情况需要处理
        
        // 左左情况
        if (balance > 1 && getBalance(node->left) >= 0) {
            return rightRotate(node);
        }
        
        // 左右情况
        if (balance > 1 && getBalance(node->left) < 0) {
            node->left = leftRotate(node->left);
            return rightRotate(node);
        }
        
        // 右右情况
        if (balance < -1 && getBalance(node->right) <= 0) {
            return leftRotate(node);
        }
        
        // 右左情况
        if (balance < -1 && getBalance(node->right) > 0) {
            node->right = rightRotate(node->right);
            return leftRotate(node);
        }
        
        return node;
    }
    
    // 查找节点的辅助函数
    bool search(Node* node, int val) {
        if (node == nullptr) {
            return false;
        }
        
        if (val == node->data) {
            return true;
        } else if (val < node->data) {
            return search(node->left, val);
        } else {
            return search(node->right, val);
        }
    }
    
    // 前序遍历辅助函数
    void preOrder(Node* node) {
        if (node != nullptr) {
            cout << node->data << " ";
            preOrder(node->left);
            preOrder(node->right);
        }
    }
    
    // 中序遍历辅助函数
    void inOrder(Node* node) {
        if (node != nullptr) {
            inOrder(node->left);
            cout << node->data << " ";
            inOrder(node->right);
        }
    }
    
    // 后序遍历辅助函数
    void postOrder(Node* node) {
        if (node != nullptr) {
            postOrder(node->left);
            postOrder(node->right);
            cout << node->data << " ";
        }
    }
    
    // 销毁树的辅助函数
    void destroyTree(Node* node) {
        if (node != nullptr) {
            destroyTree(node->left);
            destroyTree(node->right);
            delete node;
        }
    }
    
public:
    // 构造函数
    AVLTree() : root(nullptr) {}
    
    // 析构函数
    ~AVLTree() {
        destroyTree(root);
        root = nullptr;
    }
    
    // 插入节点
    void insert(int val) {
        root = insert(root, val);
    }
    
    // 删除节点
    void remove(int val) {
        if (root != nullptr) {
            root = deleteNode(root, val);
        }
    }
    
    // 查找节点
    bool search(int val) {
        return search(root, val);
    }
    
    // 前序遍历
    void preOrder() {
        cout << "前序遍历: ";
        preOrder(root);
        cout << endl;
    }
    
    // 中序遍历
    void inOrder() {
        cout << "中序遍历: ";
        inOrder(root);
        cout << endl;
    }
    
    // 后序遍历
    void postOrder() {
        cout << "后序遍历: ";
        postOrder(root);
        cout << endl;
    }
    
    // 获取树的高度
    int getHeight() {
        return height(root);
    }
};

// 测试程序
int main() {
    AVLTree tree;
    
    // 插入测试数据
    int values[] = {9, 5, 10, 0, 6, 11, -1, 1, 2};
    int n = sizeof(values) / sizeof(values[0]);
    
    for (int i = 0; i < n; i++) {
        tree.insert(values[i]);
        cout << "插入 " << values[i] << " 后，树的高度为: " << tree.getHeight() << endl;
    }
    
    // 显示各种遍历结果
    tree.preOrder();  // 预期: 9 1 0 -1 5 2 6 10 11
    tree.inOrder();   // 预期: -1 0 1 2 5 6 9 10 11
    tree.postOrder(); // 预期: -1 0 2 6 5 1 11 10 9
    
    // 查找测试
    int key = 6;
    cout << "查找 " << key << ": " << (tree.search(key) ? "找到" : "未找到") << endl;
    
    key = 3;
    cout << "查找 " << key << ": " << (tree.search(key) ? "找到" : "未找到") << endl;
    
    // 删除测试
    key = 10;
    tree.remove(key);
    cout << "删除 " << key << " 后，中序遍历: ";
    tree.inOrder();  // 预期: -1 0 1 2 5 6 9 11
    
    key = 1;
    tree.remove(key);
    cout << "删除 " << key << " 后，中序遍历: ";
    tree.inOrder();  // 预期: -1 0 2 5 6 9 11
    
    return 0;
}