C++ 伸展树与红黑树原理及实现详解

3.2 旋转操作

旋转是维护树结构的基础，分为左旋和右旋。注意更新父指针和子指针的指向关系。

template<typename K, typename V>
class RBTree {
private:
    using Node = RBNode<K, V>;
    Node* root;
    Node* nil; // 哨兵节点

    void left_rotate(Node* x) {
        Node* 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 right_rotate(Node* y) {
        Node* 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->left) y->parent->left = x;
        else y->parent->right = x;
        x->right = y;
        y->parent = x;
    }

3.3 插入与修复

插入逻辑遵循 BST 规则，但插入后需检查并修复红黑树性质。主要分两种情况处理叔叔节点的颜色：

1. 叔叔节点为红色：父节点和叔叔节点变黑，祖父节点变红，然后以祖父节点为新当前节点继续向上检查。
2. 叔叔节点为黑色：根据当前节点、父节点、祖父节点的相对位置进行旋转（单旋或双旋）并重新着色。

    void insert_fixup(Node* z) {
        while (z->parent->color == RED) {
            if (z->parent == z->parent->parent->left) {
                Node* uncle = z->parent->parent->right;
                if (uncle->color == RED) {
                    z->parent->color = BLACK;
                    uncle->color = BLACK;
                    z->parent->parent->color = RED;
                    z = z->parent->parent;
                } else {
                    if (z == z->parent->right) {
                        z = z->parent;
                        left_rotate(z);
                    }
                    z->parent->color = BLACK;
                    z->parent->parent->color = RED;
                    right_rotate(z->parent->parent);
                }
            } else {
                Node* uncle = z->parent->parent->left;
                if (uncle->color == RED) {
                    z->parent->color = BLACK;
                    uncle->color = BLACK;
                    z->parent->parent->color = RED;
                    z = z->parent->parent;
                } else {
                    if (z == z->parent->left) {
                        z = z->parent;
                        right_rotate(z);
                    }
                    z->parent->color = BLACK;
                    z->parent->parent->color = RED;
                    left_rotate(z->parent->parent);
                }
            }
        }
        root->color = BLACK;
    }

    void bst_insert(const K& key, const V& value) {
        Node* parent = nil;
        Node* curr = root;
        while (curr != nil) {
            parent = curr;
            if (key < curr->key) curr = curr->left;
            else if (key > curr->key) curr = curr->right;
            else {
                curr->value = value;
                return;
            }
        }
        Node* new_node = new Node(key, value, RED);
        new_node->parent = parent;
        new_node->left = nil;
        new_node->right = nil;
        if (parent == nil) root = new_node;
        else if (key < parent->key) parent->left = new_node;
        else parent->right = new_node;
    }

public:
    RBTree() {
        nil = new Node(K(), V(), BLACK);
        root = nil;
    }

    ~RBTree() {
        delete nil;
    }

    void insert(const K& key, const V& value) {
        bst_insert(key, value);
        // 查找刚插入的节点以便修复（实际工程中可优化返回引用）
        Node* cur = root;
        while(cur != nil && !(cur->key == key && cur->value == value)) {
             // 简单遍历找到对应节点用于修复，生产环境建议 bst_insert 返回节点指针
             if(key < cur->key) cur = cur->left;
             else if(key > cur->key) cur = cur->right;
             else break;
        }
        if(cur != nil && cur->color == RED) insert_fixup(cur);
    }

    void inorder() const {
        inorder_traversal(root);
        cout << endl;
    }

    void inorder_traversal(Node* node) const {
        if (node == nil) return;
        inorder_traversal(node->left);
        cout << "[" << node->key << ":" << node->value << "," 
             << (node->color == RED ? "红" : "黑") << "] ";
        inorder_traversal(node->right);
    }

    Node* find(const K& key) const {
        Node* curr = root;
        while (curr != nil) {
            if (key < curr->key) curr = curr->left;
            else if (key > curr->key) curr = curr->right;
            else return curr;
        }
        return nullptr;
    }
};

3.4 验证与测试

在实际使用中，我们可以通过中序遍历验证有序性，并通过检查黑高度验证平衡性。以下是一个简单的测试示例：

int main() {
    RBTree<int, string> rb_tree;
    rb_tree.insert(10, "A");
    rb_tree.insert(20, "B");
    rb_tree.insert(30, "C");
    rb_tree.insert(15, "D");
    rb_tree.insert(25, "E");
    rb_tree.insert(5, "F");

    cout << "红黑树中序遍历（键值：颜色）：" << endl;
    rb_tree.inorder();

    auto node = rb_tree.find(15);
    if (node) {
        cout << "\n查找键 15：值=" << node->value 
             << "，颜色=" << (node->color == RED ? "红" : "黑") << endl;
    }
    return 0;
}

4. 相关 OJ 题实战

掌握理论后，可以通过 LeetCode 上的经典题目巩固理解。

4.1 二叉搜索树中的插入操作

这是最基础的 BST 插入，不涉及平衡调整。

class Solution {
public:
    TreeNode* insertIntoBST(TreeNode* root, int val) {
        if (root == nullptr) return new TreeNode(val);
        TreeNode* cur = root;
        while (cur != nullptr) {
            if (val < cur->val) {
                if (cur->left == nullptr) {
                    cur->left = new TreeNode(val);
                    break;
                }
                cur = cur->left;
            } else {
                if (cur->right == nullptr) {
                    cur->right = new TreeNode(val);
                    break;
                }
                cur = cur->right;
            }
        }
        return root;
    }
};

4.2 将二叉搜索树变平衡

利用中序遍历获取有序数组，再递归构建平衡树。

class Solution {
private:
    void inorder(TreeNode* root, vector<int>& nums) {
        if (root == nullptr) return;
        inorder(root->left, nums);
        nums.push_back(root->val);
        inorder(root->right, nums);
    }

    TreeNode* build(const vector<int>& nums, int start, int end) {
        if (start > end) return nullptr;
        int mid = start + (end - start) / 2;
        TreeNode* node = new TreeNode(nums[mid]);
        node->left = build(nums, start, mid - 1);
        node->right = build(nums, mid + 1, end);
        return node;
    }

public:
    TreeNode* balanceBST(TreeNode* root) {
        vector<int> nums;
        inorder(root, nums);
        return build(nums, 0, nums.size() - 1);
    }
};

4.3 判断是否为平衡二叉树

通过后序遍历计算高度，若不平衡则返回 -1。

class Solution {
private:
    int getHeight(TreeNode* node) {
        if (node == nullptr) return 0;
        int leftHeight = getHeight(node->left);
        if (leftHeight == -1) return -1;
        int rightHeight = getHeight(node->right);
        if (rightHeight == -1) return -1;
        if (abs(leftHeight - rightHeight) > 1) return -1;
        return max(leftHeight, rightHeight) + 1;
    }

public:
    bool isBalanced(TreeNode* root) {
        return getHeight(root) != -1;
    }
};

4.4 二叉搜索树迭代器

利用中序遍历结果缓存，模拟迭代器行为。

class BSTIterator {
private:
    vector<int> inorderList;
    int idx;

    void inorder(TreeNode* node) {
        if (node == nullptr) return;
        inorder(node->left);
        inorderList.push_back(node->val);
        inorder(node->right);
    }

public:
    BSTIterator(TreeNode* root) {
        inorder(root);
        idx = 0;
    }

    int next() {
        return inorderList[idx++];
    }

    bool hasNext() {
        return idx < inorderList.size();
    }
};

4.5 二叉树中的最大路径和

使用递归计算最大贡献值，同时更新全局最大值。

class Solution {
private:
    int helper(TreeNode* node, int& maxSum) {
        if (node == nullptr) return 0;
        int leftContribution = max(helper(node->left, maxSum), 0);
        int rightContribution = max(helper(node->right, maxSum), 0);
        int currentPathSum = node->val + leftContribution + rightContribution;
        maxSum = max(maxSum, currentPathSum);
        return node->val + max(leftContribution, rightContribution);
    }

public:
    int maxPathSum(TreeNode* root) {
        int maxSum = INT_MIN;
        helper(root, maxSum);
        return maxSum;
    }
};

通过以上讲解和代码示例，相信你对伸展树和红黑树有了更深入的理解。在实际开发中，除非有特殊需求（如热点数据缓存），否则直接使用 STL 提供的容器通常是更稳妥的选择。希望这些内容能帮助你更好地掌握数据结构的核心思想。