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

3. 红⿊树的实现

3.1 红⿊树的结构

// RBTree.h
#pragma once
#include <iostream>
using namespace std;

enum Color { RED, BLACK };

template<class K, class V>
struct RBTreeNode {
    pair<K, V> _kv;
    RBTreeNode* _left;
    RBTreeNode* _right;
    RBTreeNode* _parent;
    Color _col;

    RBTreeNode(const pair<K, V>& kv)
        :_kv(kv), _left(nullptr), _right(nullptr), _parent(nullptr), _col(RED) {}
};

template<class K, class V>
class RBTree {
    typedef RBTreeNode<K, V> Node;
public:
    bool Insert(const pair<K, V>& kv);
    void InOrder();
    bool IsBalance();
    int Height();
    int Size();
    Node* Find(const K& key);
private:
    Node* _root = nullptr;
    // ... 辅助函数省略
};

3.2 红黑树的搜索

由于每一棵红黑树都是二叉搜索树，可以使用与普通二叉搜索树相同的算法进行搜索，不需要使用颜色信息。时间复杂度为 O(log n)。

3.3 红⿊树的插⼊

1. 按二叉搜索树规则插入，新结点默认为红色。
2. 若为空树，根结点设为黑色。
3. 若父结点为黑色，无需调整。
4. 若父结点为红色，违反性质，需根据叔叔结点颜色进行变色或旋转处理。

3.4 红⿊树的插⼊代码实现

bool Insert(const pair<K, V>& kv) {
    if (_root == nullptr) {
        _root = new Node(kv);
        _root->_col = BLACK;
        return true;
    }
    Node* parent = nullptr;
    Node* cur = _root;
    while (cur) {
        if (cur->_kv.first < kv.first) {
            parent = cur; cur = cur->_right;
        } else if (cur->_kv.first > kv.first) {
            parent = cur; cur = cur->_left;
        } else {
            return false;
        }
    }
    cur = new Node(kv);
    cur->_col = RED;
    if (parent->_kv.first < kv.first) parent->_right = cur;
    else parent->_left = cur;
    cur->_parent = parent;

    while (parent && parent->_col == RED) {
        Node* grandfather = parent->_parent;
        if (grandfather->_left == parent) {
            Node* uncle = grandfather->_right;
            if (uncle && uncle->_col == RED) {
                parent->_col = BLACK; uncle->_col = BLACK; grandfather->_col = RED;
                cur = grandfather; parent = cur->_parent;
            } else {
                if (cur == parent->_left) {
                    RotateR(grandfather);
                    parent->_col = BLACK; grandfather->_col = RED;
                } else {
                    RotateL(parent); RotateR(grandfather);
                    cur->_col = BLACK; grandfather->_col = RED;
                }
                break;
            }
        } else {
            Node* uncle = grandfather->_left;
            if (uncle && uncle->_col == RED) {
                parent->_col = BLACK; uncle->_col = BLACK; grandfather->_col = RED;
                cur = grandfather; parent = cur->_parent;
            } else {
                if (cur == parent->_right) {
                    RotateL(grandfather);
                    parent->_col = BLACK; grandfather->_col = RED;
                } else {
                    RotateR(parent); RotateL(grandfather);
                    cur->_col = BLACK; grandfather->_col = RED;
                }
                break;
            }
        }
    }
    _root->_col = BLACK;
    return true;
}

3.5 红⿊树的查找

Node* Find(const K& key) {
    Node* cur = _root;
    while (cur) {
        if (cur->_kv.first < key) cur = cur->_right;
        else if (cur->_kv.first > key) cur = cur->_left;
        else return cur;
    }
    return nullptr;
}

3.6 红⿊树的验证

检查性质：根为黑、无连续红结点、黑高度一致。

bool Check(Node* cur, int blackNum, const int blackNumRef) {
    if (cur == nullptr) {
        if (blackNum != blackNumRef) return false;
        return true;
    }
    if (cur->_col == RED && cur->_parent && cur->_parent->_col == RED) return false;
    if (cur->_col == BLACK) ++blackNum;
    return Check(cur->_left, blackNum, blackNumRef) && Check(cur->_right, blackNum, blackNumRef);
}

3.7 红⿊树的删除

删除逻辑类似二叉搜索树，若删除黑色结点可能破坏黑高度，需引入'双重黑色'概念并通过旋转变色消除，恢复红黑树特性。

3.8 红黑树模拟实现

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

enum Color { RED, BLACK };

template<typename K, typename V>
struct RBNode {
    K key;
    V value;
    Color color;
    RBNode* parent;
    RBNode* left;
    RBNode* right;
    RBNode(const K& k, const V& v, Color c = RED)
        :key(k), value(v), color(c), parent(nullptr), left(nullptr), right(nullptr) {}
};

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;
    }

    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;
    }

    Node* 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 curr; }
        }
        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;
        return new_node;
    }

    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);
    }

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

    ~RBTree() { delete nil; }

    void insert(const K& key, const V& value) {
        Node* new_node = bst_insert(key, value);
        if (new_node->color == RED) insert_fixup(new_node);
    }

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

    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;
    }
};

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.inorder();
    return 0;
}

4. 相关 OJ 题

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

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; }
                else cur = cur->left;
            } else {
                if (cur->right == nullptr) { cur->right = new TreeNode(val); break; }
                else 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 判断是否为平衡二叉树

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;
    }
};