Linux内核数据管理利器--红黑树

写在前面
[1. 红黑树的原理](#1. 红黑树的原理)
[2. 红黑树操作](#2. 红黑树操作)
- [2.1 红黑树的节点插入](#2.1 红黑树的节点插入)
- [2.2 红黑树的节点删除](#2.2 红黑树的节点删除)
- [2.3 红黑树的查询操作](#2.3 红黑树的查询操作)
[3. 红黑树操作实验](#3. 红黑树操作实验)
[附录A: 实验代码](#附录A: 实验代码)

写在前面

本文通过两个方面让读者可以深入理解Linux内核中红黑树RB Tree的实现以及使用，读完此文章，你可以收获：

红黑树的特性
红黑树的插入、删除、查询操作
在Linux内核代码中如何使用RB Tree库函数，这一部分通过一个实验带读者体会

1. 红黑树的原理

红黑树RB Tree是二叉树的一种，作为一种自平衡二叉树（一些情况下不是完全平衡的），它在最坏的情况下查询复杂度为\(O(logN)\)。与AVL树类似，尽管RB Tree查询效率不如AVL树（因为RB Tree左右子树高度差距最多接近两倍，而AVL树始终保持左右子树高度最多不超过1），但其插入删除效率高，适合用于大数据量且更新频繁的场景，例如内核IO调度算法。

红黑树在二叉树的基础上做了如下约束：

树种全部节点要么是黑色要么是红色
树的根节点是黑色的
叶节点（指NULL节点）颜色为黑色
红色节点之间不能相邻
一个节点的左子树和右子树高度（只统计黑色节点）相同

在介绍红黑树的操作前，我们先说明以下几点惯例：

所有节点在插入的时候都将是红色节点（不包括根节点，其插入时是黑色的），这样有一个好处是可以不违反约束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函数完成，其主要分以下三种情况，第一种是父节点为黑色，那么不需要做任何事情，插入红节点后该树仍然符合所有规则。

c 复制代码

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棵树，之后将祖父节点作为插入节点继续向上检查，下面的代码执行的正是这一步骤：

c 复制代码

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

第三种情况最为复杂，由于叔父节点不再是红色，故不能只靠染色来解决，其可分为以下四种：

插入节点为父节点的右节点，父节点为祖父节点的左节点；
插入节点为父节点的左节点，父节点为祖父节点的左节点；
插入节点为父节点的右节点，父节点为祖父节点的右节点；
插入节点为父节点的左节点，父节点为祖父节点的右节点；

在这四种中，第2种（左左）和第3种（右右）需要先进行一次染色解决规则4冲突，然后经过旋转解决染色后的规则5冲突。以左左为例，先将父节点染成黑色，祖父节点染成红色，此时不再有颜色冲突，但是规则5出现冲突，因为左子树显然多出一个黑色节点，所以接下来祖父节点右旋，将父节点作为祖父节点，这样就完成了两个恰到好处的事情：1）祖父节点位置的颜色再次变为黑色，这必然使得祖父不会破坏规则4；2）由于原祖父节点染成红色，所以即使其变成了右子树的节点也不影响规则5。下图展示了这一过程：

对于右右，其与左左区别在于使用左旋，原理可以参考左左自行推断。

对于第1种（右左）和第4种（左右），需要多增加一个旋转，使其变为左左或者右右，然后便可按照左左/右右的规则调整RB Tree，下图展示了右左的调整过程。

需要注意的是，不论是这四种中的哪种，最后操作的结果实际上都是在祖父节点和叔父节点直接新插入了红色节点，祖父节点颜色并没有改变，而且黑色节点数量也没有改变，所以在调整结束后无需继续向上检查。下面是内核中关于第三种情况的处理：

c 复制代码

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中，需要执行以下几步：

遍历RB Tree，找到新节点插入位置；
调用rb_link_node将节点链接到1找到的位置；
调用rb_insert_color调整RB Tree，使其符合规则。

2.2 红黑树的节点删除

红黑树的删除比插入操作更为复杂，其分为两个阶段，第一个阶段先删除节点，其技巧为：如果删除节点只有一个孩子或者没孩子，那么直接删除该节点，并链接父节点和孩子节点，代码如下：

c 复制代码

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

如果有两个孩子，那么选择删除节点的顺序下一个节点替换删除节点，既删除位置变到了删除节点的顺序下一个节点的原先位置，这样可以保证删除节点只有一个右子树（因为删除节点的顺序下一个节点是删除节点的右子树的最左边的叶子节点），代码如下：

c 复制代码

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为黑色，那么就要继续向上检查调整，代码如下：

c 复制代码

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染成黑色，代码如下：

c 复制代码

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的颜色，代码如下：

cstatic 复制代码

                 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的右节点变成了黑色，代码如下：

cstatic 复制代码

                 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内核中红黑树库提供的功能没有特定某一种排序方法，所以也没有给出查询接口。由于红黑树也是二叉排序树的一种，以升序为例，我们只需要按照以下流程即可进行查询操作：

python 复制代码

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: 实验代码

c 复制代码

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>