算法之红黑树

红黑树在编程应用中是一种很常见的算法，很多程序员用到了可能不知道而已。 比如大家经常用到的容器：map，set 等实现原理都采用了红黑树。

废话不多说，直接贴出代码：

#include <iostream>

enum Color { RED, BLACK };

struct Node {
    int key;
    Color color;
    Node* left;
    Node* right;
    Node* parent;

    Node(int key) : key(key), color(RED), left(nullptr), right(nullptr), parent(nullptr) {}
};

class RedBlackTree {
private:
    Node* root;

    // 红黑树的插入操作
    void insert(int key) {
        Node* newNode = new Node(key);

        // 执行普通的二叉搜索树插入操作
        Node* current = root;
        Node* parent = nullptr;
        while (current != nullptr) {
            parent = current;
            if (key < current->key)
                current = current->left;
            else
                current = current->right;
        }
        newNode->parent = parent;
        if (parent == nullptr)
            root = newNode;
        else if (key < parent->key)
            parent->left = newNode;
        else
            parent->right = newNode;

        // 修正红黑树的性质
        fixInsert(newNode);
    }

    // 修正红黑树的性质
    void fixInsert(Node* node) {
        while (node->parent != nullptr && node->parent->color == RED) {
            if (node->parent == node->parent->parent->left) {
                Node* uncle = node->parent->parent->right;
                if (uncle != nullptr && uncle->color == RED) {
                    // Case 1: 叔叔节点为红色
                    node->parent->color = BLACK;
                    uncle->color = BLACK;
                    node->parent->parent->color = RED;
                    node = node->parent->parent;
                } else {
                    if (node == node->parent->right) {
                        // Case 2: 叔叔节点为黑色，且当前节点是右子节点
                        node = node->parent;
                        rotateLeft(node);
                    }
                    // Case 3: 叔叔节点为黑色，且当前节点是左子节点
                    node->parent->color = BLACK;
                    node->parent->parent->color = RED;
                    rotateRight(node->parent->parent);
                }
            } else {
                Node* uncle = node->parent->parent->left;
                if (uncle != nullptr && uncle->color == RED) {
                    // Case 1: 叔叔节点为红色
                    node->parent->color = BLACK;
                    uncle->color = BLACK;
                    node->parent->parent->color = RED;
                    node = node->parent->parent;
                } else {
                    if (node == node->parent->left) {
                        // Case 2: 叔叔节点为黑色，且当前节点是左子节点
                        node = node->parent;
                        rotateRight(node);
                    }
                    // Case 3: 叔叔节点为黑色，且当前节点是右子节点
                    node->parent->color = BLACK;
                    node->parent->parent->color = RED;
                    rotateLeft(node->parent->parent);
                }
            }
        }
        root->color = BLACK;
    }

    // 左旋操作
    void rotateLeft(Node* node) {
        Node* rightChild = node->right;
        node->right = rightChild->left;
        if (rightChild->left != nullptr)
            rightChild->left->parent = node;
        rightChild->parent = node->parent;
        if (node->parent == nullptr)
            root = rightChild;
        else if (node == node->parent->left)
            node->parent->left = rightChild;
        else
            node->parent->right = rightChild;
        rightChild->left = node;
        node->parent = rightChild;
    }

    // 右旋操作
    void rotateRight(Node* node) {
        Node* leftChild = node->left;
        node->left = leftChild->right;
        if (leftChild->right != nullptr)
            leftChild->right->parent = node;
        leftChild->parent = node->parent;
        if (node->parent == nullptr)
            root = leftChild;
        else if (node == node->parent->left)
            node->parent->left = leftChild;
        else
            node->parent->right = leftChild;
        leftChild->right = node;
        node->parent = leftChild;
    }

    // 红黑树的查找操作
    Node* search(int key) {
        Node* current = root;
        while (current != nullptr && current->key != key) {
            if (key < current->key)
                current = current->left;
            else
                current = current->right;
        }
        return current;
    }

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

    // 公开的插入接口
    void insertKey(int key) {
        insert(key);
    }

    // 公开的查找接口
    bool searchKey(int key) {
        Node* node = search(key);
        return (node != nullptr);
    }
};

int main() {
    RedBlackTree tree;
    tree.insertKey(10);
    tree.insertKey(20);
    tree.insertKey(30);
    tree.insertKey(40);

    std::cout << "Search 20: " << (tree.searchKey(20) ? "Found" : "Not Found") << std::endl;
    std::cout << "Search 50: " << (tree.searchKey(50) ? "Found" : "Not Found") << std::endl;

    return 0;
}

这段代码定义了一个红黑树类 RedBlackTree，其中包含了插入操作 insert、修正红黑树性质的方法 fixInsert、左旋操作 rotateLeft、右旋操作 rotateRight，以及查找操作 search。在 main 函数中，我们创建了一个红黑树对象，并进行了一些插入和查找操作的示例。