C++手撕红黑树：从0到200行，拿下STL map底层核心

文章目录

[C++手撕红黑树：从0到200行，拿下STL map底层核心](#C++手撕红黑树：从0到200行，拿下STL map底层核心)
- [1. 红黑树的概念](#1. 红黑树的概念)
- - [1.1 红黑树的规则](#1.1 红黑树的规则)
- [1.2 红黑树如何确保最长路径不超过最短路径的2倍？](#1.2 红黑树如何确保最长路径不超过最短路径的2倍？)
- [1.3 红黑树的效率](#1.3 红黑树的效率)
- [2. 红黑树的实现](#2. 红黑树的实现)
- - [2.1 红黑树的结构](#2.1 红黑树的结构)
- [2.2 红黑树的插入](#2.2 红黑树的插入)
- - [2.2.1 插入的大概过程](#2.2.1 插入的大概过程)
  - [2.2.2 情况1：变色](#2.2.2 情况1：变色)
  - [2.2.3 情况2：单旋 + 变色](#2.2.3 情况2：单旋 + 变色)
  - [2.2.4 情况3：双旋 + 变色](#2.2.4 情况3：双旋 + 变色)
- [2.3 红黑树的插入代码实现](#2.3 红黑树的插入代码实现)
- [2.4 红黑树的查找](#2.4 红黑树的查找)
- [2.5 红黑树的验证](#2.5 红黑树的验证)

C++手撕红黑树：从0到200行，拿下STL map底层核心

1. 红黑树的概念

红黑树是一棵二叉搜索树，它的每个结点增加一个存储位来表示结点的颜色，可以是红色或者黑色。通过对任何一条从根到叶子的路径上各个结点的颜色进行约束，红黑树确保没有一条路径会比其他路径长出2倍，因而是接近平衡的。

1.1 红黑树的规则

每个结点不是红色就是黑色。
根结点是黑色的。
如果一个结点是红色的，则它的两个孩子结点必须是黑色的（即任意一条路径不会有连续的红色结点）。
对于任意一个结点，从该结点到其所有 NULL 结点的简单路径上，均包含相同数量的黑色结点。

说明：《算法导论》等书籍中补充了"每个叶子结点（NIL）都是黑色的"规则。这里所指的叶子结点不是传统意义上的叶子结点，而是我们说的空结点（NIL），也叫外部结点。引入 NIL 是为了准确标识所有路径，但在实现细节中通常忽略 NIL 结点，了解概念即可。

1.2 红黑树如何确保最长路径不超过最短路径的2倍？

由规则4可知，从根到 NULL 结点的每条路径都有相同数量的黑色结点。极端场景下，最短路径 一定是全为黑色结点的路径，假设最短路径长度为 bh（black height）。
由规则2和规则3可知，任意一条路径不会有连续的红色结点。极端场景下，最长路径 就是一黑一红间隔组成，那么最长路径的长度为 2 * bh。
综合红黑树的4点规则，理论上的全黑最短路径和一黑一红的最长路径并不一定在每棵红黑树中都存在。假设任意一条从根到 NULL 结点的路径长度为 h，那么 bh <= h <= 2 * bh。

1.3 红黑树的效率

假设 N 是红黑树中结点数量，h 是最短路径的长度，那么：

2^h - 1 <= N <= 2^(2*h) - 1

由此推出 h ≈ logN，即红黑树增删查改的最坏情况是走最长路径 2*logN，时间复杂度仍为 O(logN)。

红黑树的表达相对 AVL 树要抽象一些。AVL 树通过高度差直观地控制平衡，而红黑树通过4条规则的颜色约束间接实现了近似平衡。两者效率属于同一档次，但红黑树在插入相同数量的结点时旋转次数更少，因为它对平衡的控制没那么严格。

2. 红黑树的实现

2.1 红黑树的结构

cpp 复制代码

// 枚举值表示颜色
enum Colour {
    RED,
    BLACK
};

// 这里默认按 key/value 结构实现
template<class K, class V>
struct RBTreeNode {
    // 更新控制平衡需要加入 parent 指针
    pair<K, V> _kv;
    RBTreeNode<K, V>* _left;
    RBTreeNode<K, V>* _right;
    RBTreeNode<K, V>* _parent;
    Colour _col;

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

template<class K, class V>
class RBTree {
    typedef RBTreeNode<K, V> Node;
public:
    // ...
private:
    Node* _root = nullptr;
};

2.2 红黑树的插入

2.2.1 插入的大概过程

按二叉搜索树规则插入新结点。
如果是空树插入，新增结点为黑色；如果是非空树插入，新增结点必须为红色（否则会破坏规则4）。
非空树插入后，如果父亲结点是黑色，则插入结束；如果父亲结点是红色，则违反规则3，需要进一步处理。

约定：c 为当前结点（cur），p 为父亲（parent），g 为祖父（grandfather），u 为叔叔（uncle，即 p 的兄弟）。

2.2.2 情况1：变色

条件：c 为红，p 为红，g 为黑，u 存在且为红。
处理：将 p 和 u 变黑，g 变红，然后把 g 当作新的 c，继续往上更新。

无论 c 是 p 的左还是右，p 是 g 的左还是右，处理方法相同。
如果 g 是根，最后再将 g 变回黑色。

2.2.3 情况2：单旋 + 变色

条件：c 为红，p 为红，g 为黑，u 不存在或为黑。

如果 u 不存在，c 一定是新增结点。
如果 u 存在且为黑，c 一定不是新增，而是由情况1变色更新上来的。

处理：

p 是 g 的左，c 是 p 的左：以 g 为旋转点进行右单旋 ，再将 p 变黑，g 变红。
p 是 g 的右，c 是 p 的右：以 g 为旋转点进行左单旋 ，再将 p 变黑，g 变红。

2.2.4 情况3：双旋 + 变色

条件：c 为红，p 为红，g 为黑，u 不存在或为黑，且 c 与 p 的方向不一致。

处理：

p 是 g 的左，c 是 p 的右：先以 p 为旋转点进行左单旋 ，再以 g 为旋转点进行右单旋 ，最后将 c 变黑，g 变红。
p 是 g 的右，c 是 p 的左：先以 p 为旋转点进行右单旋 ，再以 g 为旋转点进行左单旋 ，最后将 c 变黑，g 变红。

2.3 红黑树的插入代码实现

cpp 复制代码

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) {
                    // 单旋
                    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) {
                // 情况1：叔叔存在且为红
                parent->_col = uncle->_col = BLACK;
                grandfather->_col = RED;
                cur = grandfather;
                parent = cur->_parent;
            } else {
                // 情况2/3：叔叔不存在或为黑
                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;
}

旋转代码与 AVL 树相同，只需调整指针，无需更新平衡因子。

2.4 红黑树的查找

按二叉搜索树逻辑实现，时间复杂度 O(logN)。

cpp 复制代码

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

2.5 红黑树的验证

不能简单地通过检查最长路径不超过最短路径2倍来验证，因为即使满足该条件，颜色规则也可能被破坏。必须检查4点规则：

根结点为黑色。
无连续红色结点。
每条路径黑色结点数量相同。

cpp 复制代码

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