红黑树

admin

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

#include <iostream>
#include <algorithm>
using namespace std;

// 定义节点颜色
enum Color { RED, BLACK };

// 红黑树节点结构
template <typename T>
struct Node {
    T data;
    Color color;
    Node *left, *right, *parent;
    
    // 构造函数
    Node(T val) : data(val), color(RED), left(nullptr), right(nullptr), parent(nullptr) {}
};

// 红黑树类
template <typename T>
class RedBlackTree {
private:
    Node<T>* root;
    Node<T>* NIL;  // 哨兵节点，代表空节点
    
    // 左旋转操作
    void leftRotate(Node<T>* x) {
        Node<T>* y = x->right;
        x->right = y->left;
        
        if (y->left != NIL)
            y->left->parent = x;
        
        y->parent = x->parent;
        
        if (x->parent == NIL)
            root = y;
        else if (x == x->parent->left)
            x->parent->left = y;
        else
            x->parent->right = y;
        
        y->left = x;
        x->parent = y;
    }
    
    // 右旋转操作
    void rightRotate(Node<T>* y) {
        Node<T>* x = y->left;
        y->left = x->right;
        
        if (x->right != NIL)
            x->right->parent = y;
        
        x->parent = y->parent;
        
        if (y->parent == NIL)
            root = x;
        else if (y == y->parent->right)
            y->parent->right = x;
        else
            y->parent->left = x;
        
        x->right = y;
        y->parent = x;
    }
    
    // 插入后修复红黑树性质
    void insertFixup(Node<T>* z) {
        while (z->parent->color == RED) {
            if (z->parent == z->parent->parent->left) {
                Node<T>* y = z->parent->parent->right;  // 叔叔节点
                
                // 情况1：叔叔是红色
                if (y->color == RED) {
                    z->parent->color = BLACK;
                    y->color = BLACK;
                    z->parent->parent->color = RED;
                    z = z->parent->parent;
                } 
                // 情况2：叔叔是黑色，且z是右孩子
                else {
                    if (z == z->parent->right) {
                        z = z->parent;
                        leftRotate(z);
                    }
                    // 情况3：叔叔是黑色，且z是左孩子
                    z->parent->color = BLACK;
                    z->parent->parent->color = RED;
                    rightRotate(z->parent->parent);
                }
            } 
            // 镜像情况：父节点是祖父节点的右孩子
            else {
                Node<T>* y = z->parent->parent->left;  // 叔叔节点
                
                // 情况1：叔叔是红色
                if (y->color == RED) {
                    z->parent->color = BLACK;
                    y->color = BLACK;
                    z->parent->parent->color = RED;
                    z = z->parent->parent;
                } 
                // 情况2：叔叔是黑色，且z是左孩子
                else {
                    if (z == z->parent->left) {
                        z = z->parent;
                        rightRotate(z);
                    }
                    // 情况3：叔叔是黑色，且z是右孩子
                    z->parent->color = BLACK;
                    z->parent->parent->color = RED;
                    leftRotate(z->parent->parent);
                }
            }
            
            if (z == root)
                break;
        }
        root->color = BLACK;  // 确保根节点是黑色
    }
    
    //  Transplant操作：用v子树替换u子树
    void transplant(Node<T>* u, Node<T>* v) {
        if (u->parent == NIL)
            root = v;
        else if (u == u->parent->left)
            u->parent->left = v;
        else
            u->parent->right = v;
        v->parent = u->parent;
    }
    
    // 查找最小节点
    Node<T>* minimum(Node<T>* node) {
        while (node->left != NIL)
            node = node->left;
        return node;
    }
    
    // 删除后修复红黑树性质
    void deleteFixup(Node<T>* x) {
        while (x != root && x->color == BLACK) {
            if (x == x->parent->left) {
                Node<T>* w = x->parent->right;  // 兄弟节点
                
                // 情况1：兄弟是红色
                if (w->color == RED) {
                    w->color = BLACK;
                    x->parent->color = RED;
                    leftRotate(x->parent);
                    w = x->parent->right;
                }
                
                // 情况2：兄弟的两个孩子都是黑色
                if (w->left->color == BLACK && w->right->color == BLACK) {
                    w->color = RED;
                    x = x->parent;
                } 
                else {
                    // 情况3：兄弟的左孩子是红色，右孩子是黑色
                    if (w->right->color == BLACK) {
                        w->left->color = BLACK;
                        w->color = RED;
                        rightRotate(w);
                        w = x->parent->right;
                    }
                    
                    // 情况4：兄弟的右孩子是红色
                    w->color = x->parent->color;
                    x->parent->color = BLACK;
                    w->right->color = BLACK;
                    leftRotate(x->parent);
                    x = root;
                }
            } 
            // 镜像情况：x是其父节点的右孩子
            else {
                Node<T>* w = x->parent->left;  // 兄弟节点
                
                // 情况1：兄弟是红色
                if (w->color == RED) {
                    w->color = BLACK;
                    x->parent->color = RED;
                    rightRotate(x->parent);
                    w = x->parent->left;
                }
                
                // 情况2：兄弟的两个孩子都是黑色
                if (w->right->color == BLACK && w->left->color == BLACK) {
                    w->color = RED;
                    x = x->parent;
                } 
                else {
                    // 情况3：兄弟的右孩子是红色，左孩子是黑色
                    if (w->left->color == BLACK) {
                        w->right->color = BLACK;
                        w->color = RED;
                        leftRotate(w);
                        w = x->parent->left;
                    }
                    
                    // 情况4：兄弟的左孩子是红色
                    w->color = x->parent->color;
                    x->parent->color = BLACK;
                    w->left->color = BLACK;
                    rightRotate(x->parent);
                    x = root;
                }
            }
        }
        x->color = BLACK;
    }
    
    // 递归删除树
    void destroyTree(Node<T>* node) {
        if (node != NIL) {
            destroyTree(node->left);
            destroyTree(node->right);
            delete node;
        }
    }
    
    // 前序遍历辅助函数
    void preOrderHelper(Node<T>* node) {
        if (node != NIL) {
            cout << node->data << "(" << (node->color == RED ? "红" : "黑") << ") ";
            preOrderHelper(node->left);
            preOrderHelper(node->right);
        }
    }
    
    // 中序遍历辅助函数
    void inOrderHelper(Node<T>* node) {
        if (node != NIL) {
            inOrderHelper(node->left);
            cout << node->data << "(" << (node->color == RED ? "红" : "黑") << ") ";
            inOrderHelper(node->right);
        }
    }
    
    // 后序遍历辅助函数
    void postOrderHelper(Node<T>* node) {
        if (node != NIL) {
            postOrderHelper(node->left);
            postOrderHelper(node->right);
            cout << node->data << "(" << (node->color == RED ? "红" : "黑") << ") ";
        }
    }
    
    // 查找辅助函数
    Node<T>* searchHelper(Node<T>* node, T val) {
        if (node == NIL || val == node->data)
            return node;
        
        if (val < node->data)
            return searchHelper(node->left, val);
        else
            return searchHelper(node->right, val);
    }

public:
    // 构造函数
    RedBlackTree() {
        NIL = new Node<T>(0);
        NIL->color = BLACK;
        NIL->left = NIL->right = NIL->parent = NIL;
        root = NIL;
    }
    
    // 析构函数
    ~RedBlackTree() {
        destroyTree(root);
        delete NIL;
    }
    
    // 插入节点
    void insert(T val) {
        Node<T>* z = new Node<T>(val);
        z->left = NIL;
        z->right = NIL;
        
        Node<T>* y = NIL;
        Node<T>* x = root;
        
        // 找到插入位置
        while (x != NIL) {
            y = x;
            if (z->data < x->data)
                x = x->left;
            else
                x = x->right;
        }
        
        z->parent = y;
        
        if (y == NIL)
            root = z;
        else if (z->data < y->data)
            y->left = z;
        else
            y->right = z;
        
        // 新节点初始颜色为红色
        z->color = RED;
        
        // 修复红黑树性质
        insertFixup(z);
    }
    
    // 删除节点
    void remove(T val) {
        Node<T>* z = searchHelper(root, val);
        if (z == NIL) {
            cout << "值 " << val << " 不在树中，无法删除" << endl;
            return;
        }
        
        Node<T>* y = z;
        Node<T>* x;
        Color yOriginalColor = y->color;
        
        // 只有右子树或没有子树
        if (z->left == NIL) {
            x = z->right;
            transplant(z, z->right);
        }
        // 只有左子树
        else if (z->right == NIL) {
            x = z->left;
            transplant(z, z->left);
        }
        // 有两个子树
        else {
            y = minimum(z->right);
            yOriginalColor = y->color;
            x = y->right;
            
            if (y->parent == z)
                x->parent = y;
            else {
                transplant(y, y->right);
                y->right = z->right;
                y->right->parent = y;
            }
            
            transplant(z, y);
            y->left = z->left;
            y->left->parent = y;
            y->color = z->color;
        }
        
        delete z;
        
        // 如果删除的是黑色节点，需要修复红黑树性质
        if (yOriginalColor == BLACK)
            deleteFixup(x);
    }
    
    // 查找节点
    bool search(T val) {
        Node<T>* result = searchHelper(root, val);
        return result != NIL;
    }
    
    // 前序遍历
    void preOrder() {
        cout << "前序遍历: ";
        preOrderHelper(root);
        cout << endl;
    }
    
    // 中序遍历
    void inOrder() {
        cout << "中序遍历: ";
        inOrderHelper(root);
        cout << endl;
    }
    
    // 后序遍历
    void postOrder() {
        cout << "后序遍历: ";
        postOrderHelper(root);
        cout << endl;
    }
    
    // 获取根节点
    Node<T>* getRoot() {
        return root;
    }
};

// 测试程序
int main() {
    RedBlackTree<int> tree;
    
    // 插入测试数据
    int values[] = {7, 3, 18, 10, 22, 8, 11, 26};
    int n = sizeof(values) / sizeof(values[0]);
    
    for (int i = 0; i < n; i++) {
        tree.insert(values[i]);
        cout << "插入 " << values[i] << " 后，中序遍历: ";
        tree.inOrder();
    }
    
    // 显示各种遍历结果
    tree.preOrder();
    tree.inOrder();
    tree.postOrder();
    
    // 查找测试
    int key = 10;
    cout << "查找 " << key << ": " << (tree.search(key) ? "找到" : "未找到") << endl;
    
    key = 15;
    cout << "查找 " << key << ": " << (tree.search(key) ? "找到" : "未找到") << endl;
    
    // 删除测试
    key = 18;
    tree.remove(key);
    cout << "删除 " << key << " 后，中序遍历: ";
    tree.inOrder();
    
    key = 3;
    tree.remove(key);
    cout << "删除 " << key << " 后，中序遍历: ";
    tree.inOrder();
    
    key = 7;
    tree.remove(key);
    cout << "删除 " << key << " 后，中序遍历: ";
    tree.inOrder();
    
    return 0;
}