【数据结构】红黑树插入自平衡实现

前言

新学期学业繁忙，以至于没什么时间没更新博客；今天整理了红黑树的相关知识，代码都是手写简化版，如果哪里有问题希望批评指正

Red Black Tree 简介

红黑树（Red Black Tree）是一种自平衡的二叉搜索树，它满足以下性质：

• 红黑树的节点为红色或黑色；

• 根节点和叶子节点都为黑色，且叶子节点为 NIL；

• 如果一个节点为红色，则它的子节点都为黑色；

• 从根节点到 NIL 节点的每条路径上的黑色节点数量相同。

R-B 树不是绝对的平衡，而是*“相对的平衡”*，即：每条路径上的黑色节点数量相同 。

额外条件：

• R-B 树的每个节点额外存储了一个 color 字段，用于确保树在插入和删除时保持平衡；

• 最长路径的长度不能超过最短路径的 2 倍。

红黑树手撕代码

先来看看 R-B 树的结构：

#include 
using namespace std;

enum Color{ RED, BLACK };

class TreeNode{
private:
    int val;
    Color color;
    TreeNode* left;
    TreeNode* right;
    TreeNode* parent;
    
    TreeNode(int val) : val(val), color(RED), left(nullptr), right(nullptr), parent(nullptr) {}
};

class RedBlackTree{
private:
    TreeNode* root;
    TreeNode* nil_leaf;  // 所有的空叶子节点共享一个节点 
  
public:
    RedBlackTree()
    {
        nil_leaf = new TreeNode(-1); // 哨兵节点的值可以是任何不会与实际数据冲突的值
        nil_leaf->color = BLACK;     // 哨兵节点一定是黑色的
        root = nil_leaf;
    }
    ~RedBlackTree(){}
    // search, insert, remove, inorder等函数
};

再来看看查找操作，这与二叉搜索树的搜索思路是相同的：

bool search(int val) {
    TreeNode* currentNode = root;
    while (currentNode) {
        if (val == currentNode->val) {
            return true; // 找到了匹配的值
        }
        else if (val < currentNode->val) {
            currentNode = currentNode->left;
        }
        else {
            currentNode = currentNode->right;
        }
    }
    return false; // 没有找到匹配的值
}

接着来看看 R-B 树的插入操作，R-B 树的插入位置和 BST 树是一样的，新节点为红色。

如果遇到以下情况，需要特殊处理：

• 如果新节点为根节点，则转为黑色；

• 如果新节点的“父”和“叔”为红色，则“爷父叔”转换颜色，且保证根节点为黑色；

• 如果新节点的“父”为红，“叔”为黑，“爷父孙”取中间值往上提为黑根，其余节点为红，左给左，右给右。

了解完插入原理，现在来实现插入的代码：

void insert(int val){
    TreeNode* node = new TreeNode(val);
    // 如果 R-B树 为空，新节点作黑根
    if(!root){
        root = node;
        node->color = BLACK;
        return;
    }
    // 如果 R-B树 非空，找到插入位置
    TreeNode* ptr = root;
    while (ptr) {
        if (node->val < ptr->val && ptr->left) {
            ptr = ptr->left;
        }
        else if (node->val < ptr->val && !ptr->left) {
            ptr->left = node; node->parent = ptr;
            break;
        }
        if (node->val > ptr->val && ptr->right) {
            ptr = ptr->right;
        }
        else if (node->val > ptr->val && !ptr->right){
            ptr->right = node; node->parent = ptr;
            break;
        }
        if (node->val == ptr->val) return;
    }
    // 如果 "父节点" 为红，则需要自平衡
    if (ptr->color == RED) balance(node);
}

到这里，我们处理了上述三种情况的第一种，现在编写 balance 函数处理另外两种情况。

void balance(TreeNode* node){
    TreeNode* parent = node->parent;
    TreeNode* grandparent = parent->parent;
    TreeNode* uncle = nullptr;
  
    if (grandparent->left == parent) {
        uncle = grandparent->right;
    } else {
        uncle = grandparent->left;
    }
    
    if(uncle->color == RED){
        // 反转颜色
        grandparent->color = RED；
        uncle->color = parent->color = BLACK;
        // 防止 grandparent 与它的 parent 同为红色
        if(grandparent->parent && grandparent->parent->color == RED){
            balance(grandparent);
        }
    }
    else{
        // 根据位置关系，进行左旋/右旋
        if (parent->left == node && grandparent->left == parent) {
            // 左左情况
            rightRotate(grandparent);
            std::swap(parent->color, grandparent->color);
        } else if (parent->right == node && grandparent->right == parent) {
            // 右右情况
            leftRotate(grandparent);
            std::swap(parent->color, grandparent->color);
        } else if (parent->left == node && grandparent->right == parent) {
            //  > 情况
            rightRotate(parent);
            leftRotate(grandparent);
            std::swap(node->color, grandparent->color);
        } else if (parent->right == node && grandparent->left == parent) {
            //  < 情况
            leftRotate(parent);
            rightRotate(grandparent);
            std::swap(node->color, grandparent->color);
        }
    }
    // 确保根节点为黑色
    root->color = BLACK;
}

从 左旋/右旋 的参数，不难推出：
左旋——将当前节点（参数）逆时针向下旋转，右子树位置上移；
右旋——将当前节点（参数）顺时针向下旋转，左子树位置上移。

是的，我们还需要编写 左旋/右旋 的代码：

void LeftRotate(TreeNode* node){
    TreeNode* parent = node->parent;
    TreeNode* rchild = node->right;
    
    node->right = rchild->left;
    if(rchild->left){
        rchild->left->parent = node;
    }
    rchild->left = node;
    rchild->parent = node->parent;
    node->parent = rchild;
    // 上升的变黑，下降的变红
    node->color = RED;
    rchild->color = BLACK;
}

void RightRotate(TreeNode* node){
    TreeNode* parent = node->parent;
    TreeNode* lchild = node->left;
    
    node->left = lchild->right;
    if(lchild->right){
        lchild->right->parent = node;
    }
    lchild->right = node;
    lchild->parent = node->parent;
    node->parent = lchild;
    // 上升的变黑，下降的变红
    node->color = RED;
    lchild->color = BLACK;
}

总结

到这里，红黑树的插入操作及自平衡原理就已经介绍完了。

其余操作有空再整理！