Linux核心資料管理利器--紅黑樹

目錄

• 寫在前面
• 1. 紅黑樹的原理
• 2. 紅黑樹操作
  • 2.1 紅黑樹的節點插入
  • 2.2 紅黑樹的節點刪除
  • 2.3 紅黑樹的查詢操作
• 3. 紅黑樹操作實驗
• 附錄A: 實驗程式碼

---

寫在前面

本文透過兩個方面讓讀者可以深入理解Linux核心中紅黑樹RB Tree的實現以及使用，讀完此文章，你可以收穫：

• 紅黑樹的特性
• 紅黑樹的插入、刪除、查詢操作
• 在Linux核心程式碼中如何使用RB Tree庫函式，這一部分透過一個實驗帶讀者體會

1. 紅黑樹的原理

紅黑樹RB Tree是二叉樹的一種，作為一種自平衡二叉樹（一些情況下不是完全平衡的），它在最壞的情況下查詢複雜度為\(O(logN)\)。與AVL樹類似，儘管RB Tree查詢效率不如AVL樹（因為RB Tree左右子樹高度差距最多接近兩倍，而AVL樹始終保持左右子樹高度最多不超過1），但其插入刪除效率高，適合用於大資料量且更新頻繁的場景，例如核心IO排程演算法。
紅黑樹在二叉樹的基礎上做了如下約束：

1. 樹種全部節點要麼是黑色要麼是紅色
2. 樹的根節點是黑色的
3. 葉節點（指NULL節點）顏色為黑色
4. 紅色節點之間不能相鄰
5. 一個節點的左子樹和右子樹高度（只統計黑色節點）相同

在介紹紅黑樹的操作前，我們先說明以下幾點慣例：

• 所有節點在插入的時候都將是紅色節點（不包括根節點，其插入時是黑色的），這樣有一個好處是可以不違反約束1，2，3和5，對於約束1，2和3是顯然的，對於5，由於新增紅色節點並不會影響其父節點及以上節點左右子樹黑色節點數量，故不違反約束5。因此，在插入節點後，只需判斷是否違反約束4。
• 一顆紅黑樹中，某一節點左右子樹節點高度差不會超過2倍，考慮一種極限情況：左子樹黑色節點高度為x，且最長路徑中不存在紅色節點，這是允許的，右子樹有黑色節點高度為x，這樣滿足約束5，除此之外，右子樹最長路徑黑色幾點之間都由紅色節點隔開（滿足約束4），故右子樹總高度為2x-1，約等於2x。

2. 紅黑樹操作

在Linux核心程式碼中僅提供了紅黑樹節點連結、索引、調整、刪除等基礎操作，不包含特定含義的查詢、插入等操作：

• void rb_insert_color(struct rb_node *, struct rb_root *);，檢查調整一個指定節點，通常與rb_link_node搭配使用；
• void rb_erase(struct rb_node *, struct rb_root *);，從樹中刪除一個指定節點；
• struct rb_node *rb_next(struct rb_node *);，返回一個節點的下一個節點（順序的）；
• struct rb_node *rb_prev(struct rb_node *);，返回一個節點的上一個節點（順序的）；
• struct rb_node *rb_first(struct rb_root *);，返回樹中的第一個節點（順序的）；
• struct rb_node *rb_last(struct rb_root *);，返回樹中的最後一個節點（順序的）；
• void rb_replace_node(struct rb_node *victim, struct rb_node *new, struct rb_root *root);，用new替換節點victim；
• inline void rb_link_node(struct rb_node * node, struct rb_node * parent, struct rb_node ** rb_link)，將一個節點連結到樹中指定位置，parent是父節點，rb_link指定了連結父節點的位置是左還是右。

2.1 紅黑樹的節點插入

根據第一個部分我們所講的內容可知，一個節點插入RB Tree時會被染成紅色，因此只需要檢查插入時是否違反規則4，既插入節點與其父節點是否都是紅色，然後做出相應的調整，這些工作由rb_insert_color函式完成，其主要分以下三種情況，第一種是父節點為黑色，那麼不需要做任何事情，插入紅節點後該樹仍然符合所有規則。

void rb_insert_color(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *parent, *gparent;

    while ((parent = node->rb_parent) && parent->rb_color == RB_RED)
    {
        ... // 檢查與處理
    }

    root->rb_node->rb_color = RB_BLACK; // 保證根節點是黑色的
}

由程式碼可知，只要父節點為黑色那麼可以直接退出。第二種情況是父節點為紅色，此時違反規則4，但是其叔父節點（父節點的父節點的另一個子節點）也是紅色，如下圖所示，左邊四個樹包含了全部這種情況，A是祖父，B是插入節點的父節點，E是插入節點。

這種情況下，可以直接將父節點和叔父節點染成黑色，祖父節點染成紅色，這樣插入節點的父節點解決了規則4，同時祖父節點左右子樹黑色節點高度仍然相同，例如上圖中的第5棵樹，之後將祖父節點作為插入節點繼續向上檢查，下面的程式碼執行的正是這一步驟：

void rb_insert_color(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *parent, *gparent;

    while ((parent = node->rb_parent) && parent->rb_color == RB_RED)
    {
        gparent = parent->rb_parent; // 祖父節點

        if (parent == gparent->rb_left)
        {
            {
                register struct rb_node *uncle = gparent->rb_right;
                if (uncle && uncle->rb_color == RB_RED)
                {
                    uncle->rb_color = RB_BLACK;
                    parent->rb_color = RB_BLACK;
                    gparent->rb_color = RB_RED;
                    node = gparent;
                    continue;
                }
            }
            ... // 其他檢查和處理
        } else {
            {
                register struct rb_node *uncle = gparent->rb_left;
                if (uncle && uncle->rb_color == RB_RED)
                {
                    uncle->rb_color = RB_BLACK;
                    parent->rb_color = RB_BLACK;
                    gparent->rb_color = RB_RED;
                    node = gparent;
                    continue;
                }
            }
            ... // 其他檢查和處理
        }
    }

    root->rb_node->rb_color = RB_BLACK;
}

第三種情況最為複雜，由於叔父節點不再是紅色，故不能只靠染色來解決，其可分為以下四種：

1. 插入節點為父節點的右節點，父節點為祖父節點的左節點；
2. 插入節點為父節點的左節點，父節點為祖父節點的左節點；
3. 插入節點為父節點的右節點，父節點為祖父節點的右節點；
4. 插入節點為父節點的左節點，父節點為祖父節點的右節點；

在這四種中，第2種（左左）和第3種（右右）需要先進行一次染色解決規則4衝突，然後經過旋轉解決染色後的規則5衝突。以左左為例，先將父節點染成黑色，祖父節點染成紅色，此時不再有顏色衝突，但是規則5出現衝突，因為左子樹顯然多出一個黑色節點，所以接下來祖父節點右旋，將父節點作為祖父節點，這樣就完成了兩個恰到好處的事情：1）祖父節點位置的顏色再次變為黑色，這必然使得祖父不會破壞規則4；2）由於原祖父節點染成紅色，所以即使其變成了右子樹的節點也不影響規則5。下圖展示了這一過程：

對於右右，其與左左區別在於使用左旋，原理可以參考左左自行推斷。
對於第1種（右左）和第4種（左右），需要多增加一個旋轉，使其變為左左或者右右，然後便可按照左左/右右的規則調整RB Tree，下圖展示了右左的調整過程。

需要注意的是，不論是這四種中的哪種，最後操作的結果實際上都是在祖父節點和叔父節點直接新插入了紅色節點，祖父節點顏色並沒有改變，而且黑色節點數量也沒有改變，所以在調整結束後無需繼續向上檢查。下面是核心中關於第三種情況的處理：

static void __rb_rotate_left(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *right = node->rb_right;

    if ((node->rb_right = right->rb_left))
        right->rb_left->rb_parent = node;
    right->rb_left = node;

    if ((right->rb_parent = node->rb_parent))
    {
        if (node == node->rb_parent->rb_left)
            node->rb_parent->rb_left = right;
        else
            node->rb_parent->rb_right = right;
    }
    else
        root->rb_node = right;
    node->rb_parent = right;
}

static void __rb_rotate_right(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *left = node->rb_left;

    if ((node->rb_left = left->rb_right))
        left->rb_right->rb_parent = node;
    left->rb_right = node;

    if ((left->rb_parent = node->rb_parent))
    {
        if (node == node->rb_parent->rb_right)
            node->rb_parent->rb_right = left;
        else
            node->rb_parent->rb_left = left;
    }
    else
        root->rb_node = left;
    node->rb_parent = left;
}

void rb_insert_color(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *parent, *gparent;

    while ((parent = node->rb_parent) && parent->rb_color == RB_RED)
    {
        gparent = parent->rb_parent;

        if (parent == gparent->rb_left)
        {
            {
                register struct rb_node *uncle = gparent->rb_right;
                ... // 叔父為紅色的處理
            }

            if (parent->rb_right == node)
            {
                register struct rb_node *tmp;
                __rb_rotate_left(parent, root);
                tmp = parent; 
                parent = node;
                node = tmp;
            }

            parent->rb_color = RB_BLACK;
            gparent->rb_color = RB_RED;
            __rb_rotate_right(gparent, root);
        } else {
            {
                register struct rb_node *uncle = gparent->rb_left;
                ... // 叔父為紅色的處理
            }

            if (parent->rb_left == node)
            {
                register struct rb_node *tmp;
                __rb_rotate_right(parent, root);
                tmp = parent;
                parent = node;
                node = tmp;
            }

            parent->rb_color = RB_BLACK;
            gparent->rb_color = RB_RED;
            __rb_rotate_left(gparent, root);
        }
    }

    root->rb_node->rb_color = RB_BLACK;
}

在Linux核心中，如果需要插入一個節點到RB Tree中，需要執行以下幾步：

1. 遍歷RB Tree，找到新節點插入位置；
2. 呼叫rb_link_node將節點連結到1找到的位置；
3. 呼叫rb_insert_color調整RB Tree，使其符合規則。

2.2 紅黑樹的節點刪除

紅黑樹的刪除比插入操作更為複雜，其分為兩個階段，第一個階段先刪除節點，其技巧為：如果刪除節點只有一個孩子或者沒孩子，那麼直接刪除該節點，並連結父節點和孩子節點，程式碼如下：

void rb_erase(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *child, *parent;
    int color;

    if (!node->rb_left)
        child = node->rb_right;
    else if (!node->rb_right)
        child = node->rb_left;
    else
    {
        ... // 有兩個孩子的操作
    }

    parent = node->rb_parent;
    color = node->rb_color;

    // 連結父節點和孩子節點
    if (child)
        child->rb_parent = parent;
    if (parent)
    {
        if (parent->rb_left == node)
            parent->rb_left = child;
        else
            parent->rb_right = child;
    }
    else
        root->rb_node = child;

 color: // 第二階段：調整
    if (color == RB_BLACK)
        __rb_erase_color(child, parent, root);
}

如果有兩個孩子，那麼選擇刪除節點的順序下一個節點替換刪除節點，既刪除位置變到了刪除節點的順序下一個節點的原先位置，這樣可以保證刪除節點只有一個右子樹（因為刪除節點的順序下一個節點是刪除節點的右子樹的最左邊的葉子節點），程式碼如下：

void rb_erase(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *child, *parent;
    int color;

    if (!node->rb_left)
        ...
    else if (!node->rb_right)
        ...
    else
    {
        struct rb_node *old = node, *left;

        node = node->rb_right;
        while ((left = node->rb_left) != NULL)
            node = left;
        // 此時 node 為 刪除節點的順序下一個節點（只有右子樹或者無孩子），old 為原刪除節點
        child = node->rb_right;
        parent = node->rb_parent;
        color = node->rb_color;

        // 連結刪除節點的順序下一個節點的孩子節點和父節點
        if (child)
            child->rb_parent = parent;
        if (parent)
        {
            if (parent->rb_left == node)
                parent->rb_left = child;
            else
                parent->rb_right = child;
        }
        else
            root->rb_node = child;

        if (node->rb_parent == old) // 由於 old 是待刪除節點，而 parent 此時指向 old，所以要將 parent 指向新的 node
            parent = node;
        // node 節點替換原刪除節點
        node->rb_parent = old->rb_parent;
        node->rb_color = old->rb_color;
        node->rb_right = old->rb_right;
        node->rb_left = old->rb_left;

        // 將新 node 連結到原刪除節點 old 的父節點上
        if (old->rb_parent)
        {
            if (old->rb_parent->rb_left == old)
                old->rb_parent->rb_left = node;
            else
                old->rb_parent->rb_right = node;
        } else
            root->rb_node = node;

        // 將新 node 連結到原刪除節點 old 的子節點上
        old->rb_left->rb_parent = node;
        if (old->rb_right) // 可能刪除的右子樹只有一個節點，刪除後變為NULL
            old->rb_right->rb_parent = node;
        goto color;
    }

 color: // 第二階段：調整
    if (color == RB_BLACK)
        __rb_erase_color(child, parent, root);
}

第二階段

當在第一階段確定了刪除節點位置（通常其只有一個子樹或者沒有子樹）後，將會檢查是否要進行調色和旋轉使得節點刪除後的RB Tree再次符合規則。我們在下面透過5種大的情況來講解這一操作。
(1) 最簡單的情況是：我們刪除的節點顏色是紅色的，這意味著節點刪除後，子樹連線到其父節點後黑色節點高度不變，因此無需調整，這點可以在rb_erase函式的最後印證，因為只有刪除節點為黑色才需要執行__rb_erase_color函式。

(2) 稍微複雜的一種情況是：我們刪除的節點B顏色是黑色，同時其父節點的另一個孩子節點C顏色也是黑色且其左右孩子節點E/F也為黑色。由於父節點A的一邊少了一個黑色節點，所以應該把另一邊的黑色節點染成紅色，這樣父節點A的左右黑色節點高度相同，而且C和E/F節點顏色不衝突。對於父節點A，如果其為紅色，那正好，將其染色為黑色，這樣以A為根的子樹高度又恢復原樣，且顏色也不會衝突；如果A為黑色，那麼就要繼續向上檢查調整，程式碼如下：

static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
                 struct rb_root *root)
{
    struct rb_node *other;

    while ((!node || node->rb_color == RB_BLACK) && node != root->rb_node)
    {
        if (parent->rb_left == node)
        {
            other = parent->rb_right;
            if (other->rb_color == RB_RED)
            {
                ...
            }
            if ((!other->rb_left ||
                 other->rb_left->rb_color == RB_BLACK)
                && (!other->rb_right ||
                other->rb_right->rb_color == RB_BLACK))
            {
                other->rb_color = RB_RED;
                node = parent;
                parent = node->rb_parent;
            }
            else
            {
                ...
            }
        }
        else
        {
            other = parent->rb_left;
            if (other->rb_color == RB_RED)
            {
                ...
            }
            if ((!other->rb_left ||
                 other->rb_left->rb_color == RB_BLACK)
                && (!other->rb_right ||
                other->rb_right->rb_color == RB_BLACK))
            {
                other->rb_color = RB_RED;
                node = parent;
                parent = node->rb_parent;
            }
            else
            {
                ...
            }
        }
    }
    if (node)
        node->rb_color = RB_BLACK;
}

下面以刪除節點為左子樹為例展示了調色過程：

(3) 我們刪除的節點B顏色是黑色的，同時其父節點A的另一個孩子節點C顏色是黑色的，而C左孩子節點E為黑色，右孩子節點F為紅色。對於這種情況，可以將父節點染色成黑色左旋/右旋使得刪除節點一側增加一個黑色節點，對於另一邊，因為C因為旋轉變成了子樹根節點，所以其應該繼承原先子樹根節點顏色。除此之外，由於C不再是子樹節點，所以少了一個黑色節點，所以要把F染成黑色，程式碼如下：

static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
                 struct rb_root *root)
{
    struct rb_node *other;

    while ((!node || node->rb_color == RB_BLACK) && node != root->rb_node)
    {
        if (parent->rb_left == node)
        {
            other = parent->rb_right;
            if (other->rb_color == RB_RED)
            {
                ...
            }
            if ((!other->rb_left ||
                 other->rb_left->rb_color == RB_BLACK)
                && (!other->rb_right ||
                other->rb_right->rb_color == RB_BLACK))
            {
                ...
            }
            else
            {
                if (!other->rb_right ||
                    other->rb_right->rb_color == RB_BLACK)
                {
                    ...
                }
                other->rb_color = parent->rb_color;
                parent->rb_color = RB_BLACK;
                if (other->rb_right)
                    other->rb_right->rb_color = RB_BLACK;
                __rb_rotate_left(parent, root);
                node = root->rb_node;
                break;
            }
        }
        else
        {
            other = parent->rb_left;
            if (other->rb_color == RB_RED)
            {
                ...
            }
            if ((!other->rb_left ||
                 other->rb_left->rb_color == RB_BLACK)
                && (!other->rb_right ||
                other->rb_right->rb_color == RB_BLACK))
            {
                ...
            }
            else
            {
                if (!other->rb_left ||
                    other->rb_left->rb_color == RB_BLACK)
                {
                    ...
                }
                other->rb_color = parent->rb_color;
                parent->rb_color = RB_BLACK;
                if (other->rb_left)
                    other->rb_left->rb_color = RB_BLACK;
                __rb_rotate_right(parent, root);
                node = root->rb_node;
                break;
            }
        }
    }
    if (node)
        node->rb_color = RB_BLACK;
}

下面以刪除節點為左子樹為例展示了調色過程：

(4) 我們刪除的節點B顏色是黑色的，同時其父節點A的另一個孩子節點C顏色是黑色的，而C左孩子節點E為紅色，右孩子節點F為黑色。對於這種情況，應該先經過染色和旋轉將其變為情況(3)。其過程為將C染成紅色右旋，這樣C原先這顆子樹左右子樹黑色節點高度不變，只是C和E顏色衝突，不過這不用擔心，按照(3)的方法，C最後變成黑色，而E變成了原先A的顏色，程式碼如下：

                 struct rb_root *root)
{
    struct rb_node *other;

    while ((!node || node->rb_color == RB_BLACK) && node != root->rb_node)
    {
        if (parent->rb_left == node)
        {
            other = parent->rb_right;
            if (other->rb_color == RB_RED)
            {
                ...
            }
            if ((!other->rb_left ||
                 other->rb_left->rb_color == RB_BLACK)
                && (!other->rb_right ||
                other->rb_right->rb_color == RB_BLACK))
            {
                ...
            }
            else
            {
                if (!other->rb_right ||
                    other->rb_right->rb_color == RB_BLACK)
                {
                    register struct rb_node *o_left;
                    if ((o_left = other->rb_left))
                        o_left->rb_color = RB_BLACK;
                    other->rb_color = RB_RED;
                    __rb_rotate_right(other, root);
                    other = parent->rb_right;
                }
                other->rb_color = parent->rb_color;
                parent->rb_color = RB_BLACK;
                if (other->rb_right)
                    other->rb_right->rb_color = RB_BLACK;
                __rb_rotate_left(parent, root);
                node = root->rb_node;
                break;
            }
        }
        else
        {
            other = parent->rb_left;
            if (other->rb_color == RB_RED)
            {
                ...
            }
            if ((!other->rb_left ||
                 other->rb_left->rb_color == RB_BLACK)
                && (!other->rb_right ||
                other->rb_right->rb_color == RB_BLACK))
            {
                ...
            }
            else
            {
                if (!other->rb_left ||
                    other->rb_left->rb_color == RB_BLACK)
                {
                    register struct rb_node *o_right;
                    if ((o_right = other->rb_right))
                        o_right->rb_color = RB_BLACK;
                    other->rb_color = RB_RED;
                    __rb_rotate_left(other, root);
                    other = parent->rb_left;
                }
                other->rb_color = parent->rb_color;
                parent->rb_color = RB_BLACK;
                if (other->rb_left)
                    other->rb_left->rb_color = RB_BLACK;
                __rb_rotate_right(parent, root);
                node = root->rb_node;
                break;
            }
        }
    }
    if (node)
        node->rb_color = RB_BLACK;
}

下面以刪除節點為左子樹為例展示了調色過程：

(5) 我們刪除的節點B顏色是黑色的，同時其父節點A的另一個孩子節點C顏色是紅色的。對於這種情況，意味著父節點A必定為黑色的，而C的E/F孩子節點為黑色的，因此我們可以透過將A染成紅色左旋/右旋，然後C染成黑色，這樣，這顆子樹黑色節點高度不變，同時刪除節點一側的子樹變成了(3)或者(4)的情況，因為經過旋轉，A的右節點變成了黑色，程式碼如下：

                 struct rb_root *root)
{
    struct rb_node *other;

    while ((!node || node->rb_color == RB_BLACK) && node != root->rb_node)
    {
        if (parent->rb_left == node)
        {
            other = parent->rb_right;
            if (other->rb_color == RB_RED)
            {
                other->rb_color = RB_BLACK;
                parent->rb_color = RB_RED;
                __rb_rotate_left(parent, root);
                other = parent->rb_right;
            }
            ...
        }
        else
        {
            other = parent->rb_left;
            if (other->rb_color == RB_RED)
            {
                other->rb_color = RB_BLACK;
                parent->rb_color = RB_RED;
                __rb_rotate_right(parent, root);
                other = parent->rb_left;
            }
            ...
        }
    }
    if (node)
        node->rb_color = RB_BLACK;
}

下面以刪除節點為左子樹為例展示了調色過程：

2.3 紅黑樹的查詢操作

Linux核心中紅黑樹庫提供的功能沒有特定某一種排序方法，所以也沒有給出查詢介面。由於紅黑樹也是二叉排序樹的一種，以升序為例，我們只需要按照以下流程即可進行查詢操作：

Query x:

node = root
while node is not null and node.value != x:
    if node.value < x：
        node = node.right
    else:
        node = node.left

Return node

3. 紅黑樹操作實驗

實驗介紹：有一種物件Item，裡面包含：1）樹節點，用於管理RB Tree；2）數值，表示了物件的實際內容；3）出現次數，由於我們希望節點隨機產生，因此可能存在重複的情況，該值用於統計相同節點的數量。我們先隨機num個Item，然後使用這些Item構建出紅黑樹。最後透過輸入要擦除的物件，我們將其從樹中刪除並顯示。

下圖時程式碼執行後的效果，每個節點列印含義為[數值，出現次數，節點顏色]，最左邊為根節點，左節點在右節點上方。

附錄A: 實驗程式碼

main.c :

#include <stdio.h>
#include <stdlib.h>
#include <time.h>

#include "rbtree.h"

typedef struct _Item
{
    int val;
    int num; // appear num
    struct rb_node node;
}Item;

static int print_num = 0;
static int print_level = 0;

Item* GenerateItem();
void DFS(struct rb_node *node);

int main()
{
    int num = 0;
    Item *item, *cur, *prev = NULL;
    struct rb_node **link;
    struct rb_root root = RB_ROOT;
    
    srand(time(NULL));

    printf("Test item num: ");
    scanf("%d", &num);
    
    print_num = 0;
    printf("Generate Item[%d]:\n", num);
    /* generate a random rb tree with [num] node */
    while (num > 0)
    {
        /* randomize a rb tree node */
        item = GenerateItem();
        if (print_num == 16)
        {
            printf("\n");
            print_num = 0;
        }
        printf("%d\t", item->val);

        /* insert a rb tree node to rb tree */
        if (!root.rb_node) // empty rb tree
        {
            root.rb_node = &(item->node);
            rb_insert_color(&(item->node), &root);
            goto next_loop;
        }
        cur = rb_entry(root.rb_node, Item, node);
        /* 1. find insert position */
        while (cur)
        {
            if (cur->val == item->val) // the same item
            {
                cur->num++;
                free(item);
                goto next_loop;
            }
            else if (cur->val > item->val)
            {
                prev = cur;
                link = &(cur->node.rb_left);
                if (cur->node.rb_left == NULL)
                {
                    break;
                }
                cur = rb_entry(cur->node.rb_left, Item, node);
            }
            else
            {
                prev = cur;
                link = &(cur->node.rb_right);
                if (cur->node.rb_right == NULL)
                {
                    break;
                }
                cur = rb_entry(cur->node.rb_right, Item, node);
            }
        }
        /* 2. link node */
        rb_link_node(&(item->node), &(prev->node), link);
        /* 3. adjust */
        rb_insert_color(&(item->node), &root);
next_loop:
        num--;
    }
    
    /* print a generated rb tree */
    print_num = 0;
    print_level = 0;
    printf("\nsort result:\n");
    DFS(root.rb_node);
    printf("\n");

    /* testing erase some rb tree node */
    printf("\nTest Erase, input node value to erase its node, or input negative value to exit\n");
    while (1)
    {
        /* get the node need to erase */
        printf(">>");
        scanf("%d", &num);
        if (num < 0)
        {
            break;
        }
        /* 1. find insert position */
        if (!root.rb_node) // empty rb tree
        {
            printf("empty tree\n");
            break;
        }
        cur = rb_entry(root.rb_node, Item, node);
        while (cur)
        {
            if (cur->val == num) // the same item
            {
                break;
            }
            else if (cur->val > num)
            {
                if (cur->node.rb_left == NULL)
                {
                    cur = NULL;
                    break;
                }
                cur = rb_entry(cur->node.rb_left, Item, node);
            }
            else
            {
                if (cur->node.rb_right == NULL)
                {
                    cur = NULL;
                    break;
                }
                cur = rb_entry(cur->node.rb_right, Item, node);
            }
        }
        /* 2. do erase function */
        if (cur)
        {
            printf("erase %d\n", num);
            rb_erase(&(cur->node), &root);
            free(cur);
            DFS(root.rb_node);
            printf("\n");
        }
        else
        {
            printf("not exist\n");
        }
        printf("===================================================================\n");
    }
    
    return 0;
}

Item* GenerateItem()
{
    Item *item = (Item*)malloc(sizeof(Item));
    
    item->val = rand() % 1000;
    item->num = 1;
    
    item->node.rb_parent = NULL;
    item->node.rb_left = NULL;
    item->node.rb_right = NULL;
    
    return item;
}

void DFS(struct rb_node *node)
{
    Item *item;
    int i;
    
    if (node)
    {
        print_level++;
        DFS(node->rb_left);
        if (print_num == 4)
        {
            printf("\n");
            print_num = 0;
        }
        item = rb_entry(node, Item, node);
        for (i = 1; i < print_level; i++)
        {
            printf("            ");
        }
        printf("[%3d,%3d,%c]\n", item->val, item->num, (item->node.rb_color == RB_RED) ? 'R' : 'B');
        print_num++;
        DFS(node->rb_right);
        print_level--;
    }
}

rbtree.h :
/*
  Red Black Trees
  (C) 1999  Andrea Arcangeli <andrea@suse.de>
  
  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

  linux/include/linux/rbtree.h

  To use rbtrees you'll have to implement your own insert and search cores.
  This will avoid us to use callbacks and to drop drammatically performances.
  I know it's not the cleaner way,  but in C (not in C++) to get
  performances and genericity...

  Some example of insert and search follows here. The search is a plain
  normal search over an ordered tree. The insert instead must be implemented
  int two steps: as first thing the code must insert the element in
  order as a red leaf in the tree, then the support library function
  rb_insert_color() must be called. Such function will do the
  not trivial work to rebalance the rbtree if necessary.

-----------------------------------------------------------------------
static inline struct page * rb_search_page_cache(struct inode * inode,
                         unsigned long offset)
{
    struct rb_node * n = inode->i_rb_page_cache.rb_node;
    struct page * page;

    while (n)
    {
        page = rb_entry(n, struct page, rb_page_cache);

        if (offset < page->offset)
            n = n->rb_left;
        else if (offset > page->offset)
            n = n->rb_right;
        else
            return page;
    }
    return NULL;
}

static inline struct page * __rb_insert_page_cache(struct inode * inode,
                           unsigned long offset,
                           struct rb_node * node)
{
    struct rb_node ** p = &inode->i_rb_page_cache.rb_node;
    struct rb_node * parent = NULL;
    struct page * page;

    while (*p)
    {
        parent = *p;
        page = rb_entry(parent, struct page, rb_page_cache);

        if (offset < page->offset)
            p = &(*p)->rb_left;
        else if (offset > page->offset)
            p = &(*p)->rb_right;
        else
            return page;
    }

    rb_link_node(node, parent, p);

    return NULL;
}

static inline struct page * rb_insert_page_cache(struct inode * inode,
                         unsigned long offset,
                         struct rb_node * node)
{
    struct page * ret;
    if ((ret = __rb_insert_page_cache(inode, offset, node)))
        goto out;
    rb_insert_color(node, &inode->i_rb_page_cache);
 out:
    return ret;
}
-----------------------------------------------------------------------
*/

#ifndef    _LINUX_RBTREE_H
#define    _LINUX_RBTREE_H

// #include <linux/kernel.h>
// #include <linux/stddef.h>
#include <stdlib.h>

#define offsetof(TYPE, MEMBER) ((size_t) &((TYPE*)0)->MEMBER)
#define container_of(ptr, type, member) ({            \
        const typeof( ((type *)0)->member ) *__mptr = (ptr);    \
        (type *)( (char *)__mptr - offsetof(type,member) );})

struct rb_node
{
    struct rb_node *rb_parent;
    int rb_color;
#define    RB_RED        0
#define    RB_BLACK    1
    struct rb_node *rb_right;
    struct rb_node *rb_left;
};

struct rb_root
{
    struct rb_node *rb_node;
};

#define RB_ROOT    (struct rb_root) { NULL, }
#define    rb_entry(ptr, type, member) container_of(ptr, type, member)

extern void rb_insert_color(struct rb_node *, struct rb_root *);
extern void rb_erase(struct rb_node *, struct rb_root *);

/* Find logical next and previous nodes in a tree */
extern struct rb_node *rb_next(struct rb_node *);
extern struct rb_node *rb_prev(struct rb_node *);
extern struct rb_node *rb_first(struct rb_root *);
extern struct rb_node *rb_last(struct rb_root *);

/* Fast replacement of a single node without remove/rebalance/add/rebalance */
extern void rb_replace_node(struct rb_node *victim, struct rb_node *new, 
                struct rb_root *root);

static inline void rb_link_node(struct rb_node * node, struct rb_node * parent,
                struct rb_node ** rb_link)
{
    node->rb_parent = parent;
    node->rb_color = RB_RED;
    node->rb_left = node->rb_right = NULL;

    *rb_link = node;
}

#endif    /* _LINUX_RBTREE_H */

rbtree.c :
/*
  Red Black Trees
  (C) 1999  Andrea Arcangeli <andrea@suse.de>
  (C) 2002  David Woodhouse <dwmw2@infradead.org>
  
  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

  linux/lib/rbtree.c
*/

// #include <linux/rbtree.h>
// #include <linux/module.h>
#include "rbtree.h"

static void __rb_rotate_left(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *right = node->rb_right;

    if ((node->rb_right = right->rb_left))
        right->rb_left->rb_parent = node;
    right->rb_left = node;

    if ((right->rb_parent = node->rb_parent))
    {
        if (node == node->rb_parent->rb_left)
            node->rb_parent->rb_left = right;
        else
            node->rb_parent->rb_right = right;
    }
    else
        root->rb_node = right;
    node->rb_parent = right;
}

static void __rb_rotate_right(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *left = node->rb_left;

    if ((node->rb_left = left->rb_right))
        left->rb_right->rb_parent = node;
    left->rb_right = node;

    if ((left->rb_parent = node->rb_parent))
    {
        if (node == node->rb_parent->rb_right)
            node->rb_parent->rb_right = left;
        else
            node->rb_parent->rb_left = left;
    }
    else
        root->rb_node = left;
    node->rb_parent = left;
}

void rb_insert_color(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *parent, *gparent;

    while ((parent = node->rb_parent) && parent->rb_color == RB_RED)
    {
        gparent = parent->rb_parent;

        if (parent == gparent->rb_left)
        {
            {
                register struct rb_node *uncle = gparent->rb_right;
                if (uncle && uncle->rb_color == RB_RED)
                {
                    uncle->rb_color = RB_BLACK;
                    parent->rb_color = RB_BLACK;
                    gparent->rb_color = RB_RED;
                    node = gparent;
                    continue;
                }
            }

            if (parent->rb_right == node)
            {
                register struct rb_node *tmp;
                __rb_rotate_left(parent, root);
                tmp = parent;
                parent = node;
                node = tmp;
            }

            parent->rb_color = RB_BLACK;
            gparent->rb_color = RB_RED;
            __rb_rotate_right(gparent, root);
        } else {
            {
                register struct rb_node *uncle = gparent->rb_left;
                if (uncle && uncle->rb_color == RB_RED)
                {
                    uncle->rb_color = RB_BLACK;
                    parent->rb_color = RB_BLACK;
                    gparent->rb_color = RB_RED;
                    node = gparent;
                    continue;
                }
            }

            if (parent->rb_left == node)
            {
                register struct rb_node *tmp;
                __rb_rotate_right(parent, root);
                tmp = parent;
                parent = node;
                node = tmp;
            }

            parent->rb_color = RB_BLACK;
            gparent->rb_color = RB_RED;
            __rb_rotate_left(gparent, root);
        }
    }

    root->rb_node->rb_color = RB_BLACK;
}

static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
                 struct rb_root *root)
{
    struct rb_node *other;

    while ((!node || node->rb_color == RB_BLACK) && node != root->rb_node)
    {
        if (parent->rb_left == node)
        {
            other = parent->rb_right;
            if (other->rb_color == RB_RED)
            {
                other->rb_color = RB_BLACK;
                parent->rb_color = RB_RED;
                __rb_rotate_left(parent, root);
                other = parent->rb_right;
            }
            if ((!other->rb_left ||
                 other->rb_left->rb_color == RB_BLACK)
                && (!other->rb_right ||
                other->rb_right->rb_color == RB_BLACK))
            {
                other->rb_color = RB_RED;
                node = parent;
                parent = node->rb_parent;
            }
            else
            {
                if (!other->rb_right ||
                    other->rb_right->rb_color == RB_BLACK)
                {
                    register struct rb_node *o_left;
                    if ((o_left = other->rb_left))
                        o_left->rb_color = RB_BLACK;
                    other->rb_color = RB_RED;
                    __rb_rotate_right(other, root);
                    other = parent->rb_right;
                }
                other->rb_color = parent->rb_color;
                parent->rb_color = RB_BLACK;
                if (other->rb_right)
                    other->rb_right->rb_color = RB_BLACK;
                __rb_rotate_left(parent, root);
                node = root->rb_node;
                break;
            }
        }
        else
        {
            other = parent->rb_left;
            if (other->rb_color == RB_RED)
            {
                other->rb_color = RB_BLACK;
                parent->rb_color = RB_RED;
                __rb_rotate_right(parent, root);
                other = parent->rb_left;
            }
            if ((!other->rb_left ||
                 other->rb_left->rb_color == RB_BLACK)
                && (!other->rb_right ||
                other->rb_right->rb_color == RB_BLACK))
            {
                other->rb_color = RB_RED;
                node = parent;
                parent = node->rb_parent;
            }
            else
            {
                if (!other->rb_left ||
                    other->rb_left->rb_color == RB_BLACK)
                {
                    register struct rb_node *o_right;
                    if ((o_right = other->rb_right))
                        o_right->rb_color = RB_BLACK;
                    other->rb_color = RB_RED;
                    __rb_rotate_left(other, root);
                    other = parent->rb_left;
                }
                other->rb_color = parent->rb_color;
                parent->rb_color = RB_BLACK;
                if (other->rb_left)
                    other->rb_left->rb_color = RB_BLACK;
                __rb_rotate_right(parent, root);
                node = root->rb_node;
                break;
            }
        }
    }
    if (node)
        node->rb_color = RB_BLACK;
}

void rb_erase(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *child, *parent;
    int color;

    if (!node->rb_left)
        child = node->rb_right;
    else if (!node->rb_right)
        child = node->rb_left;
    else
    {
        struct rb_node *old = node, *left;

        node = node->rb_right;
        while ((left = node->rb_left) != NULL)
            node = left;
        child = node->rb_right;
        parent = node->rb_parent;
        color = node->rb_color;

        if (child)
            child->rb_parent = parent;
        if (parent)
        {
            if (parent->rb_left == node)
                parent->rb_left = child;
            else
                parent->rb_right = child;
        }
        else
            root->rb_node = child;

        if (node->rb_parent == old)
            parent = node;
        node->rb_parent = old->rb_parent;
        node->rb_color = old->rb_color;
        node->rb_right = old->rb_right;
        node->rb_left = old->rb_left;

        if (old->rb_parent)
        {
            if (old->rb_parent->rb_left == old)
                old->rb_parent->rb_left = node;
            else
                old->rb_parent->rb_right = node;
        } else
            root->rb_node = node;

        old->rb_left->rb_parent = node;
        if (old->rb_right)
            old->rb_right->rb_parent = node;
        goto color;
    }

    parent = node->rb_parent;
    color = node->rb_color;

    if (child)
        child->rb_parent = parent;
    if (parent)
    {
        if (parent->rb_left == node)
            parent->rb_left = child;
        else
            parent->rb_right = child;
    }
    else
        root->rb_node = child;

 color:
    if (color == RB_BLACK)
        __rb_erase_color(child, parent, root);
}

/*
 * This function returns the first node (in sort order) of the tree.
 */
struct rb_node *rb_first(struct rb_root *root)
{
    struct rb_node  *n;

    n = root->rb_node;
    if (!n)
        return NULL;
    while (n->rb_left)
        n = n->rb_left;
    return n;
}

struct rb_node *rb_last(struct rb_root *root)
{
    struct rb_node  *n;

    n = root->rb_node;
    if (!n)
        return NULL;
    while (n->rb_right)
        n = n->rb_right;
    return n;
}

struct rb_node *rb_next(struct rb_node *node)
{
    /* If we have a right-hand child, go down and then left as far
       as we can. */
    if (node->rb_right) {
        node = node->rb_right; 
        while (node->rb_left)
            node=node->rb_left;
        return node;
    }

    /* No right-hand children.  Everything down and left is
       smaller than us, so any 'next' node must be in the general
       direction of our parent. Go up the tree; any time the
       ancestor is a right-hand child of its parent, keep going
       up. First time it's a left-hand child of its parent, said
       parent is our 'next' node. */
    while (node->rb_parent && node == node->rb_parent->rb_right)
        node = node->rb_parent;

    return node->rb_parent;
}

struct rb_node *rb_prev(struct rb_node *node)
{
    /* If we have a left-hand child, go down and then right as far
       as we can. */
    if (node->rb_left) {
        node = node->rb_left; 
        while (node->rb_right)
            node=node->rb_right;
        return node;
    }

    /* No left-hand children. Go up till we find an ancestor which
       is a right-hand child of its parent */
    while (node->rb_parent && node == node->rb_parent->rb_left)
        node = node->rb_parent;

    return node->rb_parent;
}

void rb_replace_node(struct rb_node *victim, struct rb_node *new,
             struct rb_root *root)
{
    struct rb_node *parent = victim->rb_parent;

    /* Set the surrounding nodes to point to the replacement */
    if (parent) {
        if (victim == parent->rb_left)
            parent->rb_left = new;
        else
            parent->rb_right = new;
    } else {
        root->rb_node = new;
    }
    if (victim->rb_left)
        victim->rb_left->rb_parent = new;
    if (victim->rb_right)
        victim->rb_right->rb_parent = new;

    /* Copy the pointers/colour from the victim to the replacement */
    *new = *victim;
}

---

2024-3-23 created by chegxy
<chegxy's blog website>