Linux红黑树-2

红黑树的原理理解起来并不困难，无非是左右旋转，涂色。然而要编码出一个高效简洁的实现代码却挺复杂。近日学习Linux源码，对其红黑树实现作了注释，与大家分享一下。我贴的代码是我边理解边“抄写”出来的，和源始代码可能会有少许出入，但不影响理解。本文着重注释插入删除操作，其它部分不作重点。若要了解红黑树的基础知识或是Linux红黑树的使用，请另行参考其它资料。

 

先看一下节点定义：

 

typedef struct st_rb_node {
    /**
     * 本结构体四字节对齐，因而其地址中低位两个bit永远是0。
     * Linux内核开发人员非常爱惜内存，他们用其中一个空闲的位来存储颜色信息。
     * parent_color成员实际上包含了父节点指针和自己的颜色信息。
     */
    unsigned long parent_color;
#define RB_RED   0
#define RB_BLACK 1
    struct st_rb_node *left, *right;
}__attribute__((aligned(sizeof(long)))) rb_node;

 

Linux的数据结构有个特点就是结构不包括数据，而是由数据来包括结构信息。由结构体获取数据信息是通过 CONTAINER_OF 这个宏来实现的，它利用了一些编译器的特性，有兴趣的可以参考Linux的链表源码。

 

头文件中其它主要内容如下（比较好理解，未作注释）：

 

#define rb_parent(r)    ((rb_node *)((r)->parent_color & ~3))
#define rb_color(r)     ((r)->parent_color & 1)
#define rb_is_red(r)    (!rb_color(r))
#define rb_is_black(r)  rb_color(r)
#define rb_set_red(r)   do { (r)->parent_color &= ~1; } while (0)
#define rb_set_black(r) do { (r)->parent_color |= 1; } while (0)

static inline void rb_set_parent(rb_node *node, rb_node *p) {
    node->parent_color = (node->parent_color & 3) | (unsigned long) p;
}

static inline void rb_set_color(rb_node *node, int color) {
    node->parent_color = (node->parent_color & ~1) | color;
}

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

#define RB_EMPTY_ROOT(root) ((root)->node == NULL)
#define RB_EMPTY_NODE(node) (rb_parent(node) == node)
#define RB_CLEAR_NODE(node) (rb_set_parent(node, node))

void rb_insert_color(rb_node *node, rb_root *root);
void rb_erase(rb_node *node, rb_root *root);

rb_node * rb_next(const rb_node *node);
rb_node * rb_prev(const rb_node *node);
rb_node * rb_first(const rb_root *root);
rb_node * rb_last(const rb_root *root);

void rb_replace_node(rb_node *victim, rb_node *new, rb_root *root);

static inline void rb_link_node(rb_node *node, rb_node *p, rb_node **link) {
    node->parent_color = (unsigned long) p;
    node->left = node->right = NULL;
    *link = node;
}

 

两个基本操作：左旋与右旋（未注释，仅为方便参考）

 

/**
 * 对node进行左旋。有关概念可参考WIKI等资源。
 */
static void _rb_rotate_left(rb_node *node, rb_root *root) {
    rb_node *right = node->right, *parent = rb_parent(node);

    if ((node->right = right->left))
        rb_set_parent(node->right, node);

    right->left = node;
    rb_set_parent(node, right);

    if (parent) {
        if (node == parent->left)
            parent->left = right;
        else
            parent->right = right;
    } else {
        root->node = right;
    }
    rb_set_parent(right, parent);
}

/**
 * 对node进行右旋。有关概念可参考WIKI等资源。
 */
static void _rb_rotate_right(rb_node *node, rb_root *root) {
    rb_node *left = node->left, *parent = rb_parent(node);

    if ((node->left = left->right))
        rb_set_parent(node->left, node);

    left->right = node;
    rb_set_parent(node, left);

    if (parent) {
        if (node == parent->left)
            parent->left = left;
        else
            parent->right = left;
    } else {
        root->node = left;
    }
    rb_set_parent(left, parent);
}

 

插入操作的实现（前提是节点已经插入在目标位置，下面的函数只负责修正树的性质）：

 

/**
 * node是新插入的结点，有着默认的红色。
 * 本函数检查是否有违背红黑树性质的地方，并进行纠正。
 */
void rb_insert_color(rb_node *node, rb_root *root) {
    rb_node *parent, *gp;

    /**
     * 如果有父节点且父节点是红色，进行树的调整以保证树的性质。
     */
    while ((parent = rb_parent(node)) && rb_is_red(parent)) {
        gp = rb_parent(parent);

        if (parent == gp->left) {
            /**
             * 父节点是祖父节点左子的情况。
             * "else"中的情况左右相反，这里只注释"if"里的代码。
             */
            register rb_node *uncle = gp->right;
            if (uncle && rb_is_red(uncle)) {
                /**
                 * 如果有红叔，则将红叔与红父均涂黑，并将祖父节点涂红。
                 */
                rb_set_black(parent);
                rb_set_black(uncle);
                rb_set_red(gp);
                /**
                 * 现在简单路径中黑节点个数仍然平衡，但祖父变成了红色，
                 * 我们不确定有没有造成父子均红的情况，所以需要对祖父节点进行下一轮修复。
                 */
                node = gp;
                continue;
            }

            /**
             * 现在是叔节点为空或为黑的情况。
             */
            if (node == parent->right) {
                /**
                 * 如果新节点是父节点的右子，对父节点进行左旋。
                 * 旋转后树仍然平衡，但新节点占了原父节点的位子。
                 * 这两个节点交换角色后，新的父节点是红的，其左子也是红的。
                 */
                _rb_rotate_left(parent, root);
                register rb_node *tmp = node;
                node = parent;
                parent = tmp;
            }
            /**
             * 此时父红左子红，树平衡但有连续红节点。
             *
             * 父涂黑，祖父涂红，再对祖父右旋，树即调整到合法状态。
             */
            rb_set_black(parent);
            rb_set_red(gp);
            _rb_rotate_right(gp, root);
            return;
        } else {
            register rb_node *uncle = gp->left;
            if (uncle && rb_is_red(uncle)) {
                rb_set_black(parent);
                rb_set_black(uncle);
                rb_set_red(gp);
                node = gp;
                continue;
            }

            if (node == parent->left) {
                _rb_rotate_right(parent, root);
                register rb_node *tmp = node;
                node = parent;
                parent = tmp;
            }
            rb_set_black(parent);
            rb_set_red(gp);
            _rb_rotate_left(gp, root);
            return;
        }
    }
    /**
     * 若无父节点，只需将node涂黑；
     * 若父节点为黑，插入红节点不影响树的性质。
     * 循环体后直接将node涂黑，可以同时保证以上两点。
     */
     rb_set_black(root->node);
}

 

最复杂的是删除操作，分为两部分。

删除第一步，移出节点：

 

/**
 * 删除node节点（其实只是将其与树脱离），并修正树。
 * 红黑树的节点删除相对比插入更复杂，加之Linux的源码风格也不易懂，
 * 此段代码需认真体会，可以多画画图，对各种情况标记、对照一下。
 */
void rb_erase(rb_node *node, rb_root *root) {
    rb_node *child, *parent;
    int color;

    /**
     * 这里其实分成了两个大分枝。
     * 如果node有任一空子，则记录child并跳至条件语句后的语句，具体位置参考注释；
     * 如果两子均不为空，则进入"else"分枝。
     */
    if (!node->left) {
        child = node->right;
    } else if (!node->right) {
        child = node->left;
    } else {
        /**
         * 乾坤大挪移，将要删除的节点转换成只有一个子节点或无子节点的节点。
         * 具体方法是，找到它的下一个节点，也即先找其右节点，再循环找左节点，
         * 直到没有左子节点。找到的节点一定是比node"大"且紧邻node的节点。
         * 将此找到的节点放到node位置，并保持node的原有颜色，树的性质并不会受影响。
         * 此时问题转化成了删除找到的没有左子的节点。
         * （提示：也可以找左子的右节点，循环下去，得到的将是“紧邻”的小于node的节点）
         */
        rb_node *old = node, *left;
        node = node->right;
        while ((left = node->left) != NULL)
            node = left;

        /**
         * 找到节点后，将其“大挪移”过去。
         * 首先用old的父亲或root指向新找到的节点node.
         */
        if (rb_parent(old)) {
            if (rb_parent(old)->left == old)
                rb_parent(old)->left = node;
            else
                rb_parent(old)->right = node;
        } else {
            root->node = node;
        }

        /**
         * 记录找到的node节点相关信息。对于这种删除操作，node节点移至old位置后会沿用old的颜色属性，
         * 树的性质不受影响，但原node位置将失去一个节点，我们要记录节点的相关信息以便做节点的删除工作。
         */
        child = node->right;
        parent = rb_parent(node);
        color = rb_color(node);

        if (parent == old) {
            /**
             * 参考前面while循环，如果parent==old,则找到的单子节点只能是old节点的右子。
             * 这里保存的parent可以理 解为将用作node原来子节点的亲父节点。对于old是node右子的情况，
             * node取代old后，原node的父节点还是node.
             */
            parent = node;
        } else {
            /**
             * node节点的删除处理。主要是将相应的父、子联结起来。
             * 注意我们的“乾坤大挪移”可能会比较绕，删的数据是old节点，但我们将node转移过去占了old的位置，
             * 所以对于节点结构来说，我们删除的是node原来的位置。
             */
            if (child)
                rb_set_parent(child, parent);
            parent->left = child;

            node->right = old->right;
            rb_set_parent(old->right, node);
        }

        /**
         * node移过去后继续使用old节点的颜色和父节点属性。
         * parent_color同时保存了父节点信息和自己的颜色信息。
         * 关于parent_color请参考头文件中的宏。
         */
        node->parent_color = old->parent_color;
        /**
         * 之前已将old的parent的指向old子节点的指针指向了node，
         * 并且node也已经连接了old的右子（如果node与old是子与父的关系，则node保留原右子即可）
         * 上面一句完成了node指向新父亲即设置颜色的操作，现在只剩下左子了，把它搞定！
         */
        node->left = old->left;
        rb_set_parent(old->left, node);

        /**
         * 该删的该改的都已搞定，现在要检查颜色了。
         */
        goto color;
    }

    /**
     * 这里的条件语句和goto语句实在是让人晕头转向。现在对应的是本函数最上面的if语句中前两个case。
     *
     * 如果node有任一空子，也即node无子或有单子。这里跳过了前面的“乾坤大挪移”，同时相比挪移的情况，
     * 这里的区别是node可能是右子，也可能是左子。但这里的删除操作要简单些。
     *
     * 思考：挪移与不挪移的情况，最终都是把问题转到了一个无子或只有单子的节点上，
     * 为什么不能让两种情况执行一段共同的删除代码呢？
     * 我想应该是挪移本身已经使node从原位置消失了，两者删除操作有差别，加之Linux内核对速度的狂热，
     * 就分情况写成了现在这种很绕的代码。
     */
    parent = rb_parent(node);
    color = rb_color(node);

    if (child)
        rb_set_parent(child, parent);
    if (parent) {
        if (parent->left = node)
            parent->left = child;
        else
            parent->right = child;
    } else {
        root->node = child;
    }

color:
    /**
     * color是删除节点的颜色，删除的节点的层次正好在child参数与parent参数中间。
     * 如果删除的是红节点，对树没有影响，所以只有删除的是黑节点时才修正颜色。
     */
    if (color == RB_BLACK)
        _rb_erase_color(child, parent, root);
}

 

第二步，简单移出如果破坏了树的性质，则需要修复：

 

/**
 * 前置环境：在node与parent之前，刚刚删除了一个黑色节点。现在树很可能不平衡，
 * node与parent也可能红色冲突。
 * 本函数进行树的性质的修正，以使树恢复平衡。在一些情况下问题会转移到上一层节点，
 * 则须对上一层节点进行递归检查与修正。本函数中的while循环实际上实现了这种递归。
 *
 * 提示：这是红黑树里最绕的地方，看不明白可以多画画图。
 */
static void _rb_erase_color(rb_node *node, rb_node *parent, rb_root *root) {
    /**
     * other用来保存兄弟节点，这是树的修正过程中一个重要的参考节点。
     */
    rb_node *other;
    /**
     * 循环条件：node不是红节点且node不是根节点。
     * 解释：对于红节点或根节点，直接涂黑即可解决问题。
     */
    while ((!node || rb_is_black(node)) && node != root->node) {
        /**
         * 为方便程序处理各种情况，这里和节点插入一样将问题分成了左右对称的两类。
         * if-else两个分枝里代码逻辑完全相同，只是左右相反。所以我们只研究node是parent左子的情况。
         *
         * 在开始之前，我们先总结一下当前状态：
         * 1:因为删除的是黑色节点，所以node与parent都有可能是红色节点。
         * 2:node与parent之间少了一个黑色节点，则所有通过node的路径都少了一个黑色节点，不妨画图时用-1标出来；
         * 但node的兄弟节点（node一定有兄弟，可以根据删除前树的平衡性质来反推）高度并未变化，可以记作0。
         *
         * 提示：在进行旋转、涂色等操作时，可以画图观查平衡状态的变化。
         */
        if (parent->left == node) {
            other = parent->right;
            if (rb_is_red(other)) {
                /**
                 * 如果兄弟节点是红色，则父节点是黑色，交换父、兄节点的颜色，并对父节点进行左旋。
                 * 旋转后，兄节点占了老的父节点位置，且和老的父节点颜色相同，所以不会向上造成颜色冲突。
                 * 我们仍然以老的父节点为父节点来看，现在的状态是：
                 * 父节点右子保持平衡，只有经过node的路径少了一个黑色节点。
                 * 现在问题和之前相似，但node有了一个黑色的兄弟（还有一个红色父亲）。
                 */
                rb_set_black(other);
                rb_set_red(parent);
                _rb_rotate_left(parent, root);
                /**
                 * other指向新的兄弟节点。other现在必然是一个黑色节点，而不会是空。这一点可以根据旋转之前树的性质反证。
                 */
                other = parent->right;
            }
            /**
             * 此时状态：
             * node有黑色兄弟，父亲可能黑也可能红。
             */
            if ((!other->left || rb_is_black(other->left)) &&
                (!other->right || rb_is_black(other->right))) {
                /**
                 * 如果other没有红色子节点，那我们就可以把other涂红，并向上转移问题。
                 * other涂红的后果是，other分枝少了一个黑节点，与node分枝保持了平衡，但parent整体少了一个黑色节点。
                 * 细心的人可能会发现，如果父亲是红色的，父亲与兄弟有颜色冲突，直接向上转移能纠正吗？当然是可以的。
                 * 向上一层之后，while里的黑节点判断失败，会直接执行while后面的语句，直接将parent涂黑，则树恢复平衡。
                 */
                rb_set_red(other);
                node = parent;
                parent = rb_parent(node);
            } else {
                /**
                 * 现在黑兄弟有红子节点，父亲颜色未知。
                 */
                if (!other->right || rb_is_black(other->right)) {
                    /**
                     * 如果黑兄弟右节点为空或为黑，则左节点一定是红的，我们想办法把它调整为右子为红。
                     * 至于为什么，看后面就知了。
                     * other->left与other交换颜色，对other进行右旋，other指向新的兄弟。根据右旋转的特点，
                     * 可知现在other仍然是黑色，且它有了一个红色右子。同时other分枝高度不变，颜色也没有冲突。
                     */
                    rb_set_black(other->left);
                    rb_set_red(other);
                    _rb_rotate_right(other, root);
                    other = parent->right;
                }
                /**
                 * 此时状态：黑兄弟有红色右子节点。
                 * 不管parent是什么颜色，把other涂成父亲的颜色（之后旋转，other占据父亲的位置，向上没有颜色冲突），
                 * 把父亲涂黑，把黑兄的other涂黑，这时node分枝高度可能有变化也可能没变化，other分枝多了一个黑节点。
                 * 现在对父亲进行左旋转。旋转后的情况是右边分枝（原other右子）少了一个黑节点，重归平衡；
                 * 左边分枝则增加了一个黑节点，也恢复了平衡。此时也没有颜色冲突
                 */
                rb_set_color(other, rb_color(parent));
                rb_set_black(parent);
                rb_set_black(other->right);
                _rb_rotate_left(parent, root);
                /**
                 * 树已平衡，node置为根节点，并跳出循环。
                 * 疑问：这里为什么还要检查根节点呢？
                 */
                node = root->node;
                break;
            }
        } else {
            other = parent->left;
            if (rb_is_red(other)) {
                rb_set_black(other);
                rb_set_red(parent);
                _rb_rotate_right(parent, root);
                other = parent->left;
            }
            if ((!other->left || rb_is_black(other->left)) &&
                (!other->right || rb_is_black(other->right))) {
                rb_set_red(other);
                node = parent;
                parent = rb_parent(node);
            } else {
                if (!other->left || rb_is_black(other->left)) {
                    rb_set_black(other->right);
                    rb_set_red(other);
                    _rb_rotate_left(other, root);
                    other = parent->left;
                }
                rb_set_color(other, rb_color(parent));
                rb_set_black(parent);
                rb_set_black(other->left);
                _rb_rotate_right(parent, root);
                node = root->node;
                break;
            }
        }
    }
    /**
     * 对于红节点或根节点，直接涂黑即可解决问题。
     */
    if (node)
        rb_set_black(node);
}

参考地址：