紅黑樹之插入實現

紅黑樹

性質

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