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

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

红黑树概述

红黑树本质上是一棵二叉搜索树，每个节点额外增加了一个颜色位（红色或黑色）。通过对路径上颜色的约束，它确保了没有一条路径会比其他路径长出两倍，从而实现了近似平衡。

核心规则

1. 颜色：每个节点非红即黑。
2. 根节点：必须是黑色。
3. 红色限制：若节点为红色，其子节点必须为黑色（无连续红节点）。
4. 黑高一致：从任一节点到其所有叶子（NIL）的简单路径上，包含相同数量的黑色节点。

注：《算法导论》中提到的'叶子节点（NIL）都是黑色的'是为了方便理论推导。在实际代码实现中，通常不需要显式创建 NIL 节点，只需在逻辑上处理空指针即可。

为什么最长路径不超过最短路径的 2 倍？

根据规则 4，每条路径的黑色节点数量（黑高 bh）是固定的。极端情况下，最短路径全为黑色，长度为 bh。而根据规则 2 和 3，任意路径不能有连续红节点，因此最长路径是一黑一红交替，长度最多为 2 * bh。综合来看，树的高度 h 满足 bh <= h <= 2 * bh。

效率分析

假设节点数为 N，高度为 h，则有 2^h - 1 <= N <= 2^(2*h) - 1。由此可推导出 h ≈ logN。这意味着红黑树的增删查改操作在最坏情况下的时间复杂度仍为 O(logN)。

相比 AVL 树，红黑树对平衡的控制更宽松，插入时旋转次数更少，因此在频繁插入的场景下表现更优，这也是 STL 容器选择它作为底层结构的原因。

红黑树的实现细节

数据结构定义

我们需要一个节点结构体，包含键值对、左右子节点、父节点以及颜色属性。为了便于回溯调整，父指针必不可少。

enum Colour { RED, BLACK };

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

    RBTreeNode(const std::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:
    // ... 成员函数声明
private:
    Node* _root = nullptr;
};

插入逻辑详解

插入新节点时，我们先按二叉搜索树规则找到位置。如果是空树，新节点设为黑色；否则设为红色（设为黑色会破坏规则 4）。如果父节点是黑色，无需调整；如果父节点也是红色，则违反了规则 3，需要修复。

这里约定：cur 为当前节点，p 为父亲，g 为祖父，u 为叔叔（父亲的兄弟）。

情况 1：变色

条件：cur 红，p 红，g 黑，u 存在且为红。 处理：将 p 和 u 变黑，g 变红。此时 g 可能违反规则 2（若 g 是根），或者 g 与新的父节点冲突，因此将 g 视为新的 cur 继续向上检查。若 g 是根，最后将其变回黑色。

情况 2：单旋 + 变色

条件：cur 红，p 红，g 黑，u 不存在或为黑。 处理：

• 若 p 是 g 左，cur 是 p 左：以 g 为轴右旋，p 变黑，g 变红。
• 若 p 是 g 右，cur 是 p 右：以 g 为轴左旋，p 变黑，g 变红。

情况 3：双旋 + 变色

条件：cur 红，p 红，g 黑，u 不存在或为黑，且 cur 与 p 方向不一致（如 p 是 g 左，cur 是 p 右）。 处理：先以 p 为轴旋转一次，使结构变为情况 2，再以 g 为轴旋转一次，最后调整颜色。

插入代码实现

下面是完整的插入逻辑，包含了查找、插入及三种情况的修复。注意旋转代码与 AVL 树类似，但红黑树不需要维护平衡因子。

bool Insert(const std::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;
}

查找与验证

查找操作遵循二叉搜索树的标准逻辑，时间复杂度 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) {
            std::cout << "存在黑色结点数量不相等的路径" << std::endl;
            return false;
        }
        return true;
    }
    if (root->_col == RED && root->_parent->_col == RED) {
        std::cout << root->_kv.first << "存在连续的红色结点" << std::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);
}

通过上述步骤，我们就完成了一棵功能完备的红黑树实现。理解这些旋转和变色逻辑，对于掌握 STL 容器乃至整个 C++ 标准库的设计思想都至关重要。