红黑树之插入实现

红黑树

性质

1. 红黑树的结点都是红色或者黑色
2. 根结点是黑色
3. 所有叶子都是黑色（这里的叶子结点是空结点）
4. 每个红色结点必须有两个黑色的子结点
5. 从任何一个节点到其每个叶子的所有简单路径都包含相同数目的黑色结点

---

• 性质1和性质3总是能够保持着；
• 性质4只有在这些情况下才会发生作用：
  • 增加红色结点
  • 将黑色结点重新绘制成红色结点
  • 旋转
• 性质5在这些情况下才会发生作用：
  • 增加黑色结点
  • 将红色结点重新绘制黑色结点
  • 旋转

举例：

插入

用BST的方法将结点插入，将该结点标记为红色的（因为如果标记为黑色，则会导致根结点到叶子结点的路径会多出一个黑结点，无法满足性质5，而且不容易进行调整），插入的情况包括下面几种：

1. 插入到一个空的树，插入结点则为根结点，只需要将红色结点重新转染成黑色结点来满足性质2；
2. 新结点的父结点为黑色，满足所有条件；
3. 新结点的父结点为红色，因为性质2和性质4，所以树必然有祖父结点，则又包括以下的情况：

   1. 父亲结点和叔父结点均为红色，显然无法满足性质4，则将父亲结点和叔父结点绘制成黑色，祖父结点设置成红色，但是仍然无法满足情况，比如考虑到祖父结点可能是根结点，则无法满足性质2，或者祖父结点的父结点是红色的，则无法满足性质4，这时需要将祖父结点作为新的结点来看待进行各种情况的判断，涉及到对祖父结点的递归；

      

   2. 父亲结点为红色同时叔父结点为黑色或者从缺，这里又分为两种情况，新插入结点为父亲结点的左子结点和右子结点（假设其中父亲结点为祖父结点的左子结点），区别在于旋转的方向，显然，这棵树父亲结点既然为红色，那么其祖父结点则为黑色（性质4），不然无法满足前提。

      1. 新插入结点为父亲结点的左子结点，那么就构成了一个左左的情况，在之前平衡树中提到过，如果要将其进行平衡，则需要对父结点进行一次单右旋转，形成一个父亲结点为相对根结点，子结点和祖父结点为子结点的树，同时将父亲结点的红色改为黑色，祖父结点更改为红色，这下之前无法满足的性质4和性质5就满足了；

         

      2. 新插入结点为父亲结点的右子结点，那么就会构成一个左右的情况，在之前的平衡树也提到过要进行一次双旋转，先对新结点进行一次单左旋转，变成了左左的结构，再进行一次单右旋转，从而达到满足所有性质；

         

   3. 父亲结点是祖父结点的右结点，参考平衡树进行相应的操作，原理是一致的

实现

自然先看头文件，如下：

typedef enum {
    RB_RED = 0,
    RB_BLACK
} RBColor;

struct RBTreeNode {
    RBColor rb_color;
    int rb_key;
    struct RBTreeNode *rb_left;
    struct RBTreeNode *rb_right;
    struct RBTreeNode *rb_parent;

    RBTreeNode(int key) : rb_key(key), rb_color(RB_BLACK), rb_left(nullptr), rb_right(nullptr), rb_parent(nullptr) {}
};

struct RBTreeRoot {
    struct RBTreeNode* rb_node;
};

class RBTree {
public:
    RBTree();
    ~RBTree();

    void insert(int val);
    void print();
private:
    void _left_rotate(struct RBTreeRoot *root, struct RBTreeNode* node);
    void _right_rotate(struct RBTreeRoot *root, struct RBTreeNode* node);
    RBTreeNode* insert_node(struct RBTreeNode* node);
    void _insert_node(struct RBTreeNode* node);
    void _print(struct RBTreeNode* root);

    struct RBTreeRoot* _root;
};

当然，为了使用更方便，还定义了一些红定义和内联函数；

#define rb_parent(r)    ((struct RBTreeNode *)((r)->rb_parent))
#define rb_color(r)     ((r)->rb_color)
#define rb_is_red(r)    ((r)->rb_color & RB_RED)
#define rb_is_black(r)  ((r)->rb_color & RB_BLACK)
#define rb_set_red(r)   ((r)->rb_color = RB_RED)
#define rb_set_black(r) ((r)->rb_color = RB_BLACK)

static inline void rb_set_reds(int count, ...) {
    va_list args;
    va_start(args, count);
    while (count--) {
        rb_set_red(va_arg(args, struct RBTreeNode*));
    }
    va_end(args);
}

static inline void rb_set_blacks(int count, ...) {
    va_list args;
    va_start(args, count);
    while (count--) {
        rb_set_black(va_arg(args, struct RBTreeNode*));
    }
    va_end(args);
}

static inline void rb_set_parent(struct RBTreeNode* cb, struct RBTreeNode* p) {
    cb->rb_parent = p;
}

static inline void rb_set_left(struct RBTreeNode* cb, struct RBTreeNode* p) {
    cb->rb_left = p;
    rb_set_parent(p, cb);
}

static inline void rb_set_right(struct RBTreeNode* cb, struct RBTreeNode* p) {
    cb->rb_right = p;
    rb_set_parent(p, cb);
}

static inline bool rb_is_left(struct RBTreeNode* cb, struct RBTreeNode* p) {
    return p->rb_left == cb;
}

static inline bool rb_is_right(struct RBTreeNode* cb, struct RBTreeNode* p) {
    return p->rb_right == cb;
}

真正的实现在这里，其操作可以参考平衡二叉树：

RBTree::RBTree() {
    _root = new RBTreeRoot();
}

RBTree::~RBTree() {
    delete _root;
}

/*
 * 对红黑树的节点(x)进行左旋转
 *
 * 左旋示意图(对节点x进行左旋)：
 *      px                              px
 *     /                               /
 *    x                               y
 *   /        --(左旋)-->           /          #
 *  lx   y                          x  ry
 *     /                          /  
 *    ly   ry                     lx  ly
 *
 *
 */
void RBTree::_left_rotate(struct RBTreeRoot *root, struct RBTreeNode *node) {
    struct RBTreeNode *right = node->rb_right, *parent = rb_parent(node);

    // 第一步：将ly连接到x的右结点上
    rb_set_right(node, right->rb_left);

    // 第二步：将x设置为y的左子结点
    rb_set_left(right, node);

    // 第三步：将y设置为px的子结点
    if (parent) {
        if (rb_is_left(node, parent)) {
            rb_set_left(right, parent);
        }
        else {
            rb_set_right(right, parent);
        }
    }
    else {
        root->rb_node = right; // 根结点
    }
}

/*
 * 对红黑树的节点(y)进行右旋转
 *
 * 右旋示意图(对节点y进行左旋)：
 *            py                               py
 *           /                                /
 *          y                                x
 *         /        --(右旋)-->            /                       #
 *        x   ry                           lx   y
 *       /                                    /                    #
 *      lx  rx                                rx  ry
 *
 */
void RBTree::_right_rotate(struct RBTreeRoot *root, struct RBTreeNode *node) {
    struct RBTreeNode *left = node->rb_left, *parent = rb_parent(node);

    // 第一步：将rx设置为y的左子结点
    rb_set_left(node, left->rb_right);

    // 第二步：将y设置为x的右子结点
    rb_set_right(left, node);

    // 第三步：将x设置为py的子结点
    if (parent) {
        if (rb_is_left(node, parent)) {
            rb_set_left(parent, left);
        }
        else {
            rb_set_right(parent, left);
        }
    }
}

void RBTree::_insert_node(struct RBTreeNode *node) {
    struct RBTreeNode *parent, *g_parent;

    // 满足性质4
    while ((parent = rb_parent(node)) && rb_is_red(parent)) {
        g_parent = rb_parent(parent);

        if (rb_is_left(parent, g_parent)) {
            {
                // case 1：叔叔结点是红色
                // 寄存器变量，提高效率
                struct RBTreeNode *uncle = g_parent->rb_right;
                // 无法满足性质4
                if (uncle && rb_is_red(uncle)) {
                    // step1：将父亲和叔叔结点设置成黑色
                    rb_set_blacks(2, parent, uncle);
                    // step2：将祖父设置成红色（因为之前必然为黑色，不然无法满足性质4）
                    rb_set_red(g_parent);
                    // step3：递归检查祖父结点
                    node = g_parent;
                    continue;
                }
            }

            // 无法满足性质5
            // case 2：叔叔结点是黑色，并且当前结点在右边，必然要进行双旋转
            if (rb_is_right(node, parent)) {
                struct RBTreeNode *temp;

                // step 1：将父亲结点进行左旋
                _left_rotate(_root, parent); // 此时父结点为当前结点的左子结点
                // step 2：将当前结点和父结点进行交换
                temp = parent;
                parent = node;
                node = temp;
            }

            // 此时父亲结点和当前结点均是红色，无法满足性质4和性质5
            // case 3：叔叔结点是黑色，并且当前结点在左边，只用单旋转
            // step 1：将父亲结点改成改成黑色，祖父结点改成红色，以便后面进行旋转后，
            // 红色的左子结点和祖父结点为黑色的父结点的子结点
            rb_set_black(parent);
            rb_set_red(g_parent);
            // step 2：右旋转
            _right_rotate(_root, g_parent);  // 经过右旋转后，红色均分布在两边
        }
        else {  // 顺序相反而已
            {
                // case 4：叔叔结点是红色
                // 寄存器变量，提高效率
                struct RBTreeNode *uncle = g_parent->rb_left;
                // 无法满足性质4
                if (uncle && rb_is_red(uncle)) {
                    // step1：将父亲和叔叔结点设置成黑色
                    rb_set_blacks(2, parent, uncle);
                    // step2：将祖父设置成红色（因为之前必然为黑色，不然无法满足性质4）
                    rb_set_red(g_parent);
                    // step3：递归检查祖父结点
                    node = g_parent;
                    continue;
                }
            }

            // 无法满足性质5
            // case 5：叔叔结点是黑色，并且当前结点在左边，必然要进行双旋转
            if (rb_is_left(node, parent)) {
                struct RBTreeNode *temp;

                // step 1：将父亲结点进行左旋
                _right_rotate(_root, parent); // 此时父结点为当前结点的右子结点
                // step 2：将当前结点和父结点进行交换
                temp = parent;
                parent = node;
                node = temp;
            }

            // 此时父亲结点和当前结点均是红色，无法满足性质4和性质5
            // case 3：叔叔结点是黑色，并且当前结点在右边，只用单旋转
            // step 1：将父亲结点改成改成黑色，祖父结点改成红色，以便后面进行旋转后，
            // 红色的左子结点和祖父结点为黑色的父结点的子结点
            rb_set_black(parent);
            rb_set_red(g_parent);
            // step 2：左旋转
            _left_rotate(_root, g_parent);  // 经过左旋转后，红色均分布在两边
        }
    }
}

// 参照BST的插入方法
RBTreeNode* RBTree::insert_node(struct RBTreeNode *node) {
    struct RBTreeNode* temp = _root->rb_node;
    struct RBTreeNode* temp_parent = nullptr;

    while (temp != nullptr) {
        temp_parent = temp;
        if (node->rb_key < temp->rb_key) {
            temp = temp->rb_left;
        }
        else {
            temp = temp->rb_right;
        }
    }

    // 设置子结点
    if (temp_parent != nullptr) {
        if (node->rb_key < temp_parent->rb_key) {
            rb_set_left(temp_parent, node);
        }
        else {
            rb_set_right(temp_parent, node);
        }
    }
    else {
        _root->rb_node = node;  // 根结点
    }

    return node;
}

void RBTree::insert(int val) {
    struct RBTreeNode* node = new RBTreeNode(val);
    node = this->insert_node(node);
    this->_insert_node(node);
}

void RBTree::_print(struct RBTreeNode* root) {
    if (root->rb_left) {
        _print(root->rb_left);
    }
    std::cout << root->rb_key << "   ";
    if (root->rb_right) {
        _print(root->rb_right);
    }
}

void RBTree::print() {
    _print(_root->rb_node);
}