C++ 实现红黑树：深入理解 STL map 底层原理

插入操作

插入新节点时遵循 BST 规则，但需注意颜色处理：

1. 若为空树，新节点设为黑色作为根。
2. 若非空树，新节点必须设为红色。若设为黑色会破坏规则 4（改变某条路径的黑高）。
3. 如果父节点是黑色，无需调整；如果是红色，则违反规则 3，需修复。

这里约定：cur 为当前节点，parent 为父亲，grandfather 为祖父，uncle 为叔叔。

情况一：变色

当叔叔节点存在且为红色时，将父母变黑，祖父变红，然后以祖父为新当前节点继续向上检查。若祖父是根，最后将其变回黑色。

情况二：单旋 + 变色

当叔叔不存在或为黑色，且当前节点与父亲方向一致时（左左或右右），进行单旋并交换颜色。

• 左左：祖父右旋，父变黑，祖父变红。
• 右右：祖父左旋，父变黑，祖父变红。

情况三：双旋 + 变色

当叔叔不存在或为黑色，且当前节点与父亲方向不一致时（左右或右左），先对父亲单旋，再对祖父单旋，最后变色。

• 左右：父左旋，祖父右旋，当前变黑，祖父变红。
• 右左：父右旋，祖父左旋，当前变黑，祖父变红。

代码实现

下面是插入的核心逻辑，注意旋转函数与 AVL 树类似，但只需调整指针和颜色。

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 (parent == grandfather->_left) {
            Node* uncle = grandfather->_right;
            if (uncle && uncle->_col == RED) {
                // 情况 1：叔叔存在且为红
                parent->_col = uncle->_col = BLACK;
                grandfather->_col = RED;
                cur = grandfather;
                parent = cur->_parent;
            } else {
                // 情况 2/3：叔叔不存在或为黑
                if (cur == parent->_left) {
                    // 情况 2：单旋
                    RotateR(grandfather);
                    parent->_col = BLACK;
                    grandfather->_col = RED;
                } else {
                    // 情况 3：双旋
                    RotateL(parent);
                    RotateR(grandfather);
                    cur->_col = BLACK;
                    grandfather->_col = RED;
                }
                break;
            }
        } else {
            Node* uncle = grandfather->_left;
            if (uncle && uncle->_col == RED) {
                // 情况 1：叔叔存在且为红
                parent->_col = uncle->_col = BLACK;
                grandfather->_col = RED;
                cur = grandfather;
                parent = cur->_parent;
            } else {
                // 情况 2/3：叔叔不存在或为黑
                if (cur == parent->_right) {
                    // 情况 2：单旋
                    RotateL(grandfather);
                    parent->_col = BLACK;
                    grandfather->_col = RED;
                } else {
                    // 情况 3：双旋
                    RotateR(parent);
                    RotateL(grandfather);
                    cur->_col = BLACK;
                    grandfather->_col = RED;
                }
                break;
            }
        }
    }
    _root->_col = BLACK; // 确保根节点为黑
    return true;
}

查找与验证

查找逻辑与普通 BST 一致，O(logN)。验证则需要递归检查三条规则：根为黑、无连续红、黑高一致。

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

bool Check(Node* root, int blackNum, const int refNum) {
    if (root == nullptr) {
        if (refNum != blackNum) {
            cout << "存在黑色结点数量不相等的路径" << endl;
            return false;
        }
        return true;
    }
    if (root->_col == RED && root->_parent->_col == RED) {
        cout << root->_kv.first << "存在连续的红色结点" << endl;
        return false;
    }
    if (root->_col == BLACK)
        blackNum++;
    return Check(root->_left, blackNum, refNum) && Check(root->_right, blackNum, refNum);
}

bool IsBalance() {
    if (_root == nullptr) return true;
    if (_root->_col == RED) return false;

    int refNum = 0;
    Node* cur = _root;
    while (cur) {
        if (cur->_col == BLACK) ++refNum;
        cur = cur->_left;
    }
    return Check(_root, 0, refNum);
}