C++ 红黑树插入与平衡调整详解

前言

红黑树是二叉搜索树的一种变体，基于红黑树实现的结构包括 map 和 set。它是一种近似平衡的二叉搜索树，通过引入颜色属性来判断简单路径的长度，从而达到平衡的目的。

红黑树引入了以下规则：

1. 节点不是红色就是黑色。
2. 根节点是黑色的。
3. 不存在连续的两个红色节点。
4. 每条简单路径上的黑色节点数目应该相同。

只要满足以上条件，最短路径长度不超过最长路径的两倍。本文将着重介绍插入部分的调整以及判断整个树是否满足红黑树的条件。

节点定义

红黑树有颜色属性，使用枚举表示 RED 和 BLACK。采用三叉链结构，包含 parent 指针。使用 key-value 结构存储数据。

enum Colour { BLACK, RED };

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

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

初始化为红色是为了方便插入调整。整体红黑树类结构与 AVL 树类似。

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

插入部分

插入逻辑与二叉搜索树类似，区别在于插入后需要调整颜色。如果插入的是根节点，颜色必须设为黑色。

插入流程

1. 查找位置：按照二叉搜索树规则找到插入位置。
2. 连接节点：将新节点连接到父节点下，新节点默认颜色为红色。
3. 颜色调整：检查是否违反红黑树性质，若违反则进行变色或旋转。

bool Insert(const pair<T, V>& kv) {
    if (_root == nullptr) {
        _root = new Node(kv);
        _root->_col = BLACK;
        return true;
    }

    Node* root = _root;
    Node* parent = nullptr;

    // 查找插入位置
    while (root) {
        if (kv.first > root->_kv.first) {
            parent = root;
            root = root->_right;
        } else if (kv.first < root->_kv.first) {
            parent = root;
            root = root->_left;
        } else {
            return false; // 已存在
        }
    }

    Node* cur = new Node(kv);
    
    // 连接节点
    if (parent->_kv.first > kv.first) {
        parent->_left = cur;
    } else {
        parent->_right = 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) {
                parent->_col = uncle->_col = BLACK;
                grandfather->_col = RED;
                cur = grandfather;
                parent = grandfather->_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 = uncle->_col = BLACK;
                grandfather->_col = RED;
                cur = grandfather;
                parent = grandfather->_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;
}

判断部分

判断红黑树是否达标，主要检查四条规则。根节点必须为黑色。检查是否存在连续红色节点。检查每条路径黑色节点数是否一致。

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

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

private:
bool Check(Node* root, int blackNum, const int base) {
    if (root == nullptr) {
        if (base != blackNum) {
            return false;
        }
        return true;
    }

    if (root->_col == RED && root->_parent->_col == RED) {
        return false;
    }

    if (root->_col == BLACK) {
        blackNum++;
    }

    return Check(root->_left, blackNum, base) && Check(root->_right, blackNum, base);
}

相对于 AVL 树来说，红黑树在保持平衡的同时减少了旋转次数，性能通常更优。