C++ 红黑树的设计原理与插入实现详解

红黑树的实现

结构定义

在书写平衡树的代码之前，通常需要先完善它的结构。红黑树是一个典型的三叉链，节点需要包含父节点以及左右孩子节点。由于红黑树是通过颜色进行平衡的，还需要设置好一个枚举来确定节点的颜色，默认颜色为红色。红黑树通过 Key 和 Value 存放数据，Key 作为标记，Value 为核心数据。

enum Colour { RED, BLACK };

template<class K, class V>
struct RBTreeNode {
    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), _col(RED) {}
};

查找

红黑树的查找和二叉搜索树的查找是一样的，代码类似，效率都是 log(N)。

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

红黑树的插入

红黑树的插入比较复杂，涉及旋转和变色操作来维持规则。假设新增的节点叫 c(cur)，父亲为 p(parent)，祖父为 g(grandfather)，叔叔为 u(uncle)。

插入大致过程

插入一个值是按照二叉搜索树的规则进行的。如果是空树插入，插入的必然是黑色节点；如果是非空树插入，插入的节点一定是红色节点，因为第四条规则约束着，如果插入的是黑色节点，会破坏黑色节点数量平衡。如果父节点是红色的，违反了第三条规则。此时 c 是红色，p 也是红色，g 必然是黑色的。关键变化取决于 u 的变化，分为四种情况讨论。

bool Insert(const pair<K, V>& kv) {
    if (_root == nullptr) {
        _root = new Node(kv);
        _root->_col = BLACK; // 根节点一定是黑色的
    }
    Node* parent = nullptr;
    Node* cur = _root;
    while (cur) {
        if (cur->_kv.first > kv.first) {
            parent = cur;
            cur = cur->_left;
        } else if (cur->_kv.first < kv.first) {
            parent = cur;
            cur = cur->_right;
        } 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;
    // 开始进行变色处理
    // ... 后续调整逻辑
}

情况一：变色

如果 c 是红的，p 为红，g 为黑，u 存在并且是红色的，需要进行变色处理，即将 p 和 u 都变为黑色，g 变红。再把 g 当做新的 c，继续往上更新。如果 g 是整棵树的根，最后需要将 g 变为黑色。

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

情况二：单旋 + 变色

如果 c 是红色，p 是红色，g 为黑，u 不存在或者 u 是黑色的时候，不能单纯依靠变色修正。如果 p 是 g 的左子树，c 是 p 的左子树，以 g 为旋转点进行一次右单旋，然后让 g 变为红色，p 变为黑色。反之亦然。

if (uncle == nullptr || uncle->_col == BLACK) {
    if (parent->_left == cur) {
        RotateR(grandfather);
        parent->_col = BLACK;
        grandfather->_col = RED;
        break;
    }
} else {
    if (parent->_right == cur) {
        RotateL(grandfather);
        parent->_col = BLACK;
        grandfather->_col = RED;
        break;
    }
}

情况三：双旋 + 变色

如果 p 是 g 的左，c 是 p 的右，需要通过双旋进行调整。先以 p 为旋转点进行一次左旋，再以 g 进行一次右旋，之后让 c 变为黑色，g 变为红色。

// 左右双旋转
RotateL(parent);
RotateR(grandfather);
cur->_col = BLACK;
guardfather->_col = RED;

红黑树的验证

验证红黑树需要检验四条规则：

1. 节点非红即黑。
2. 根节点是黑色。
3. 红色节点的孩子必须是黑色。
4. 任意路径黑色节点数量相等。

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

红黑树的删除

红黑树的删除比插入更为复杂，涉及更多的旋转与变色逻辑，此处不再赘述。

完整代码

#pragma once
#include<iostream>
#include<vector>
#include<assert.h>
#include<ctime>
using namespace std;

enum Colour { RED, BLACK };

template<class K, class V>
struct RBTreeNode {
    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), _col(RED) {}
};

template<class K, class V>
class RBTree {
    typedef RBTreeNode<K, V> Node;
private:
    void RotateR(Node* parent) {
        Node* subl = parent->_left;
        Node* sublR = subl->_right;
        parent->_left = sublR;
        if (sublR) {
            sublR->_parent = parent;
        }
        Node* ppnode = parent->_parent;
        subl->_right = parent;
        parent->_parent = subl;
        if (ppnode == nullptr) {
            _root = subl;
            subl->_parent = ppnode;
        } else {
            if (subl->_kv.first > ppnode->_kv.first) {
                ppnode->_right = subl;
            } else {
                ppnode->_left = subl;
            }
            subl->_parent = ppnode;
        }
    }

    void RotateL(Node* parent) {
        Node* subL = parent->_right;
        Node* subLL = subL->_left;
        Node* ppnode = parent->_parent;
        parent->_right = subLL;
        if (subLL) {
            subLL->_parent = parent;
        }
        subL->_left = parent;
        parent->_parent = subL;
        if (ppnode == nullptr) {
            _root = subL;
            subL->_parent = nullptr;
        } else {
            if (subL->_kv.first > ppnode->_kv.first) {
                ppnode->_right = subL;
            } else {
                ppnode->_left = subL;
            }
            subL->_parent = ppnode;
        }
    }

    void _InOrder(Node* root) {
        if (root == nullptr) return;
        _InOrder(root->_left);
        cout << root->_kv.first << ":" << root->_kv.second << endl;
        _InOrder(root->_right);
    }

public:
    bool Insert(const pair<K, V>& kv) {
        if (_root == nullptr) {
            _root = new Node(kv);
            _root->_col = BLACK;
        }
        Node* parent = nullptr;
        Node* cur = _root;
        while (cur) {
            if (cur->_kv.first > kv.first) {
                parent = cur;
                cur = cur->_left;
            } else if (cur->_kv.first < kv.first) {
                parent = cur;
                cur = cur->_right;
            } 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) {
                    grandfather->_col = RED;
                    parent->_col = BLACK;
                    uncle->_col = BLACK;
                    cur = grandfather;
                    parent = cur->_parent;
                } else {
                    if (parent->_left == cur) {
                        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) {
                    grandfather->_col = RED;
                    parent->_col = BLACK;
                    uncle->_col = BLACK;
                    cur = grandfather;
                    parent = cur->_parent;
                } else {
                    if (parent->_right == cur) {
                        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;
    }

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

    void InOrder() {
        _InOrder(_root);
    }

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

    int Height() {
        return _Height(_root);
    }

    int Size() {
        return _Size(_root);
    }

protected:
    int _Size(Node* root) {
        if (root == nullptr) return 0;
        return _Size(root->_left) + _Size(root->_right) + 1;
    }

    int _Height(Node* root) {
        if (root == nullptr) return 0;
        int leftHeight = _Height(root->_left);
        int rightHeight = _Height(root->_right);
        return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
    }

private:
    Node* _root = nullptr;
};

小结

本文详细介绍了 C++ 红黑树的设计原理、基本规则、插入操作的三种调整情况及验证方法，并提供了完整的类实现代码。红黑树作为一种高效的自平衡二叉搜索树，广泛应用于标准库中。