[算法]树算法之红黑树

树算法之红黑树   

罗朝辉 (http://www.cnblogs.com/kesalin/)

本文遵循“署名-非商业用途-保持一致”创作公用协议

 

红黑树本质是二叉查找树的一种，它的性能高于普通的二叉查找树，即使是在最坏的情况下也能保证时间复杂度为O(lgn)。红黑树在每个结点上增加一个存储位表示结点的颜色（或红或黑，故称红黑树）。通过对任何一条从根到叶子的路径上各个结点着色方式的限制，红黑树可以保证没有一条路径会比其他路径长出两倍，因而是接近平衡的。

红黑树的每个结点至少包含五个域：color，key，left，right 和 parent（一般我们都会在结点中存储额外的数据 data，但前面的五个域是必不可少的），如果某个结点没有子结点或者结节点，则将相应的指针设置为空值（NIL，注意不是 NULL，NIL是一个特定的空结点对象，类似于Obj-C 中 Nil对象）。我们将这些 NIL 当作叶子结点（在实际处理过程中，往往将最底层的孩子结点和根结点的父亲都指向同一个 NIL 结点，以便于处理红黑树代码中的边界条件），而将其它结点当作内结点。

满足如下 5 个性质的二叉树就是一颗红黑树：
一，每个结点只有一种颜色，或红色或黑色；
二，根结点是黑色的；
三，每个叶结点是黑色的；
四，如果一个结点是红色的，那么它的两个孩子结点必须是黑色的；
五，对每个结点，从该结点到其子孙结点的所有路径上包含相同数目 H 的黑色结点， H 即为该节点的高度。

红黑树的实现，linux 内核中有现成实现算法，园子里也有不少同学（那谁 等）实现过。在这里我重复制造一下轮子，当然我这里的实现与前面提到的有些不同，这里是按照《算法导论》上的算法描述实现的。具体算法过程，请参考《算法导论》中红黑树那一章。

本文源码：

https://github.com/kesalin/CppSnippet/blob/master/Algorithms/RedBlackTree.cpp

下面先来看看头文件中包含的数据结构，接口函数：

#ifndef __RED_BLACK_TREE_H__
#define __RED_BLACK_TREE_H__

enum RBColor{
    RB_Red,
    RB_Black,
};

struct RBNode 
{
    RBColor color;
    int key;
    RBNode* leftChild;
    RBNode* rightChild;
    RBNode* parent;
};

typedef RBNode* RBTree;

RBNode* RBTree_nil();

// 最小关键字元素
RBNode* RBTree_minimum(RBNode* node);

// 最大关键字元素
RBNode* RBTree_maximum(RBNode* node);

// 中序遍历中的前驱
RBNode* RBTree_predecessor(RBNode* node);

// 中序遍历中的后继
RBNode* RBTree_successor(RBNode* node);

void RBTree_insert(RBTree* tree, RBNode* node);

RBNode* RBTree_delete(RBTree* tree, RBNode* node);

RBNode* RBTree_search(const RBTree tree, int key);

void RBTree_print(RBTree tree, int her = 1);

#endif    // __RED_BLACK_TREE_H__

下面来看具体实现：

#include "RedBlackTree.h"
#include <stdlib.h>
#include <stdio.h>
#include <assert.h>

static RBNode RBNode_Nil = {RB_Black, 0, 0, 0, 0};

RBNode* RBTree_nil()
{
    return &RBNode_Nil;
}

void RBTree_print(RBTree tree, int her)
{
    int i;
    RBNode* node = tree;

    assert(node);

    if (node != &RBNode_Nil) {
        for (i = 0; i < her; i++) {
            printf(" ");
        }
        printf("第 %d 层， %d(%c)/n",
            her, node->key, node->color == RB_Black ? 'B' : 'R');

        if (node->leftChild != &RBNode_Nil) {
            RBTree_print(node->leftChild, her + 1);
        }

        if (node->rightChild != &RBNode_Nil) {
            RBTree_print(node->rightChild, her + 1);
        }
    }
}

// 最小关键字元素
RBNode* RBTree_minimum(RBNode* node)
{
    assert(node);

    RBNode* temp = node;
    while (temp->leftChild != &RBNode_Nil) {
        temp = temp->leftChild;
    }

    return temp;
}


// 最大关键字元素
RBNode* RBTree_maximum(RBNode* node)
{
    assert(node);

    RBNode* temp = node;
    while (temp->rightChild != &RBNode_Nil) {
        temp = temp->rightChild;
    }

    return temp;
}

// 中序遍历中的前驱
RBNode* RBTree_predecessor(RBNode* node)
{
    assert(node);

    RBNode* child = node->leftChild;

    // 没有左孩子，返回自身
    if (child == &RBNode_Nil) {
        return node;
    }

    // 只有左孩子，则左孩子是其直接前驱
    else if (child->rightChild == &RBNode_Nil) {
        return child->leftChild;
    }

    // 左右孩子均有，则右孩子树中最大的元素为其直接前驱
    else {
        return RBTree_maximum(child->rightChild);
    }
}

// 中序遍历中的后继
RBNode* RBTree_successor(RBNode* node)
{
    // 有右孩子，则右孩子树中最小的元素为其直接后继
    if (node->rightChild != &RBNode_Nil) {
        return RBTree_minimum(node->rightChild);
    }

    // 没有右孩子，向上找到的第一个左分支节点为其直接后继，
    // 即 node 为其直接后继的左孩子树中的最大元素。
    RBNode* temp = node;
    RBNode* parent = temp->parent;

    while (parent != &RBNode_Nil && temp == parent->rightChild) {
        temp = parent;
        parent = temp->parent;
    }

    return parent;
}

RBNode* RBTree_search(const RBTree tree, int key)
{
    RBNode* node = tree;

    while (node != &RBNode_Nil) {
        if (node->key == key) {
            return node;
        }

        else if (node->key < key) {
            node = node->rightChild;
        }
        else {
            node = node->leftChild;
        }
    }

    return &RBNode_Nil;
}

// 左旋
//            node                        right
//           /    /                      /     /
//          a    right     -->         node     c
//              /     /               /    /
//             b       c             a      b
//
void RBTree_left_rotate(RBTree* tree, RBNode* node)
{
    assert(node->rightChild && (*tree)->parent == &RBNode_Nil);

    RBNode* right = node->rightChild;

    // set b
    node->rightChild = right->leftChild;
    if (right->leftChild != &RBNode_Nil) {
        right->leftChild->parent = node;
    }

    right->parent = node->parent;
    if (node->parent == &RBNode_Nil) {
        *tree = right;
    }
    else if (node->parent->leftChild == node) {
        node->parent->leftChild = right;
    }
    else {
        node->parent->rightChild = right;
    }

    right->leftChild = node;
    node->parent = right;
}

// 右旋
//            node                  left
//           /    /                /    /
//         left    c     -->      a     node
//        /     /                      /    /
//       a       b                    b      c
//
void RBTree_right_rotate(RBTree* tree, RBNode* node)
{
    assert(node->leftChild && (*tree)->parent == &RBNode_Nil);

    RBNode* left = node->leftChild;

    // set b
    node->leftChild = left->rightChild;
    if (left->rightChild != &RBNode_Nil) {
        left->rightChild->parent = node;
    }

    left->parent = node->parent;
    if (node->parent == &RBNode_Nil) {
        *tree = left;
    }
    else if (node->parent->leftChild == node) {
        node->parent->leftChild = left;
    }
    else {
        node->parent->rightChild = left;
    }

    left->rightChild = node;
    node->parent = left;
}

// 插入调整
void RBTree_insert_fixup(RBTree* tree, RBNode* node)
{
    assert(tree && node);

    RBNode* parent = NULL;
    RBNode* uncle = NULL;
    RBNode* grand = NULL;
    RBNode* temp = NULL;

    parent = node->parent;
    while (parent->color == RB_Red) {
        // 根据红黑树性质，以及 node 的父亲的颜色为红色，
        // 可以肯定 node 的祖父节点一定存在
        grand = parent->parent;

        // 确定叔父节点
        if (parent == grand->leftChild) {
            uncle = grand->rightChild;

            // case 1: 叔父节点为红色
            //         grand(B)        new node  ->    grand(R)
            //         /     /                         /      /
            //   parent(R)    uncle(R)    -->     node(B)   uncle(B)
            //   /     /      /  /                /   /        /  /
             //  a    node(R) d    e          parent  node(R)  d    e
            //       /   /                          /   /
            //      b     c                        b     c
            //
            if (uncle->color == RB_Red) {
                parent->color = RB_Black;
                uncle->color = RB_Black;
                grand->color = RB_Red;
                node = grand;
                parent = node->parent;
            }

            // case 2, case 3：叔父节点为黑色
            //         case 2     --->    case 3         -->  done
            //                      parent is as new node
            //         grand(B)          grand(B)            node(B)
            //         /     /           /      /           /      /
             //   parent(R)    d       node(R)   d      parent(R)  grand(R)
            //   /     /               /     /           /  /      /   /
            //  a    node(R)      parent(R)   c         a    b    c     d
             //       /   /         /  /
            //      b     c       a    b
            //
            else {
                // 将 case 2 装换成 case 3
                if (parent->rightChild == node) {
                    RBTree_left_rotate(tree, parent);
                    temp = parent;
                    parent = node;
                    node = temp;
                }

                // case 3
                parent->color = RB_Black;
                grand->color = RB_Red;

                RBTree_right_rotate(tree, grand);
            }
        }
        else {
            // 与上面的情况对称
            uncle = grand->leftChild;

            if (uncle->color == RB_Red) {
                parent->color = RB_Black;
                uncle->color = RB_Black;
                grand->color = RB_Red;
                node = grand;
                parent = node->parent;
            }

            else {
                // 将 case 2 装换成 case 3
                if (parent->leftChild == node) {
                    RBTree_right_rotate(tree, parent);
                    temp = parent;
                    parent = node;
                    node = temp;
                }

                // case 3
                parent->color = RB_Black;
                grand->color = RB_Red;

                RBTree_left_rotate(tree, grand);
            }
        }
    }

    (*tree)->color = RB_Black;
}

// 将节点 node 插入树 tree 内，然后将 node 着色为红色，此时，树可能不再
// 满足红黑树性质，因此调用 RBTree_insert_fixup 来对节点重新着色调整。
void RBTree_insert(RBTree* tree, RBNode* node)
{
    assert(tree && node);

    RBNode* parent = &RBNode_Nil;
    RBNode* temp = *tree;

    // 像二叉树一样，在树中查找适当的位置插入
    while (temp != &RBNode_Nil) {
        parent = temp;

        if (node->key < temp->key) {
            temp = temp->leftChild;
        }
        else {
            temp = temp->rightChild;
        }
    }

    node->parent = parent;

    // 树为空
    if (parent == &RBNode_Nil) {
        *tree = node;
    }
    else if (node->key < parent->key) {
        parent->leftChild = node;
    }
    else {
        parent->rightChild = node;
    }

    // 为节点着色
    node->leftChild = &RBNode_Nil;
    node->rightChild = &RBNode_Nil;
    node->color = RB_Red;

    // 调整树使之满足红黑树性质
    RBTree_insert_fixup(tree, node);
}

// 删除调整
void RBTree_delete_fixup(RBTree* tree, RBNode* node)
{
    RBNode* brother = NULL;
    RBNode* parent = NULL;

    while (node != *tree && node->color == RB_Black) {
        parent = node->parent;

        // 确定兄弟节点
        if (node == parent->leftChild) {
            brother = parent->rightChild;

            // case 1: 兄弟节点为红色
            if (brother->color == RB_Red) {
                brother->color = RB_Black;
                parent->color = RB_Red;

                RBTree_left_rotate(tree, parent);

                brother = node->parent->rightChild;
            }

            // case 2: 兄弟节点的两孩子均为黑色
            if (brother->leftChild->color == RB_Black
                && brother->rightChild->color == RB_Black) {
                    brother->color = RB_Red;
                    node = parent;
            }

            else {
                // case 3: 兄弟节点的左孩子为红色，右孩子为黑色
                if (brother->rightChild->color == RB_Black) {
                    brother->leftChild->color = RB_Black;
                    brother->color = RB_Red;

                    RBTree_right_rotate(tree, brother);

                    brother = node->parent->rightChild;
                }

                // case 4:兄弟节点的右孩子为红色
                brother->color = parent->color;
                parent->color = RB_Black;
                brother->rightChild->color = RB_Black;

                RBTree_left_rotate(tree, parent);

                node = *tree;
            }
        }
        else {
            brother = node->parent->leftChild;

            // case 1: 兄弟节点为红色
            if (brother->color == RB_Red) {
                brother->color = RB_Black;
                parent->color = RB_Red;

                RBTree_right_rotate(tree, parent);

                brother = node->parent->leftChild;
            }

            // case 2: 兄弟节点的两孩子均为黑色
            if (brother->leftChild->color == RB_Black
                && brother->rightChild->color == RB_Black) {
                    brother->color = RB_Red;
                    node = parent;
            }

            else {
                // case 3: 兄弟节点的左孩子为红色，右孩子为黑色
                if (brother->rightChild->color == RB_Black) {
                    brother->leftChild->color = RB_Black;
                    brother->color = RB_Red;

                    RBTree_left_rotate(tree, brother);

                    brother = node->parent->rightChild;
                }

                // case 4:兄弟节点的右孩子为红色
                brother->color = parent->color;
                parent->color = RB_Black;
                brother->leftChild->color = RB_Black;

                RBTree_right_rotate(tree, parent);

                node = *tree;
            }
        }
    }

    node->color = RB_Black;
}

// 删除
RBNode* RBTree_delete(RBTree* tree, RBNode* node)
{
    RBNode* successor = NULL;
    RBNode* temp = NULL;

    if (node->leftChild == &RBNode_Nil || node->rightChild == &RBNode_Nil) {
        successor = node;
    }
    else {
        successor = RBTree_successor(node);
    }

    if (successor->leftChild != &RBNode_Nil) {
        temp = successor->leftChild;
    }
    else {
        temp = successor->rightChild;
    }

    temp->parent = successor->parent;

    if (successor->parent == &RBNode_Nil) {
        *tree = temp;
    }
    else {
        if (successor == successor->parent->leftChild) {
            successor->parent->leftChild = temp;
        }
        else {
            successor->parent->rightChild = temp;
        }
    }

    if (successor != node) {
        node->key = successor->key;
    }

    if (successor->color == RB_Black) {
        RBTree_delete_fixup(tree, temp);
    }

    return successor;
}

测试代码，测试数组中的数字与测试步骤是经过仔细挑选的，以确保各个分支都可以测试到：

void test_redblacktree_delete(RBTree* tree, int key)
{
    RBNode* node = RBTree_search(*tree, key);
    assert(node != RBTree_nil());

    printf("/n删除节点 %d /n", node->key);
    
    node = RBTree_delete(tree, node);
    free(node);
    
    RBTree_print(*tree);
}

void test_redblacktree()
{
    const int length = 14;
    int array[length] = {
        2, 3, 4, 6, 7, 11, 9, 18, 12, 14, 19, 17, 22, 20
    };

    int i;
    RBTree tree = RBTree_nil();
    RBNode* node = NULL;

    // 插入节点生成树
    for (i = 0; i < length; i++) {
        node = (RBNode*)malloc(sizeof(RBNode));
        node->key = array[i];
        node->color = RB_Red;
        node->parent = RBTree_nil();
        node->leftChild = RBTree_nil();
        node->rightChild = RBTree_nil();

        RBTree_insert(&tree, node);    
    }

    RBTree_print(tree);

    // 插入测试
    node = (RBNode*)malloc(sizeof(RBNode));
    node->key = 21;

    printf("/n插入节点 %d/n", node->key);

    RBTree_insert(&tree, node);
    RBTree_print(tree);

    // 查找测试
    i = 6;
    node = RBTree_search(tree, i);

    if (node != RBTree_nil()) {
        printf("/n在红黑树中找到节点 %d/n", node->key);
    }
    else {
        printf("/n在红黑树中找不到节点 %d/n", i);
    }

    // 删除测试
    // 
    i = 4;// 取值 1, 2, 3, 4，分别对应 case 1, 2, 3, 4

    switch (i)
    {
    case 1:    // 兄弟为红色
        test_redblacktree_delete(&tree, 3);
        break;

    case 2:    // 兄弟为黑色，且兄弟的两孩子均为黑色
        test_redblacktree_delete(&tree, 12);
        break;

    case 3:    // 兄弟为黑色，且兄弟的左孩子为红色，右孩子均为黑色
        test_redblacktree_delete(&tree, 19);
        break;

    case 4:    // 兄弟为黑色，且兄弟的右孩子为红色
        test_redblacktree_delete(&tree, 9);
        break;
    }

    test_redblacktree_delete(&tree, 21);

    // 删除树
    for (i = 0; i < length; i++) {
        node = RBTree_search(tree, array[i]);

        if (node != RBTree_nil()) {
            printf("删除 %d/n", node->key);
            node = RBTree_delete(&tree, node);
            free(node);
        }
    }

    assert(tree == RBTree_nil());
}

测试结果
===============================================

创建红黑树
-------------------------------
 第 1 层， 6(B)
  第 2 层， 3(B)
   第 3 层， 2(B)
   第 3 层， 4(B)
  第 2 层， 12(B)
   第 3 层， 9(R)
    第 4 层， 7(B)
    第 4 层， 11(B)
   第 3 层， 18(R)
    第 4 层， 14(B)
     第 5 层， 17(R)
    第 4 层， 20(B)
     第 5 层， 19(R)
     第 5 层， 22(R)

插入节点 21
-------------------------------
 第 1 层， 6(B)
  第 2 层， 3(B)
   第 3 层， 2(B)
   第 3 层， 4(B)
  第 2 层， 12(R)
   第 3 层， 9(B)
    第 4 层， 7(B)
    第 4 层， 11(B)
   第 3 层， 18(B)
    第 4 层， 14(B)
     第 5 层， 17(R)
    第 4 层， 20(R)
     第 5 层， 19(B)
     第 5 层， 22(B)
      第 6 层， 21(R)

在红黑树中找到节点 6
-------------------------------

删除节点 9
-------------------------------
 第 1 层， 6(B)
  第 2 层， 3(B)
   第 3 层， 2(B)
   第 3 层， 4(B)
  第 2 层， 18(R)
   第 3 层， 12(B)
    第 4 层， 11(B)
     第 5 层， 7(R)
    第 4 层， 14(B)
     第 5 层， 17(R)
   第 3 层， 20(B)
    第 4 层， 19(B)
    第 4 层， 22(B)
     第 5 层， 21(R)

删除节点 21
-------------------------------
 第 1 层， 6(B)
  第 2 层， 3(B)
   第 3 层， 2(B)
   第 3 层， 4(B)
  第 2 层， 18(R)
   第 3 层， 12(B)
    第 4 层， 11(B)
     第 5 层， 7(R)
    第 4 层， 14(B)
     第 5 层， 17(R)
   第 3 层， 20(B)
    第 4 层， 19(B)
    第 4 层， 22(B)

删除 2
删除 3
删除 4
删除 6
删除 7
删除 11
删除 18
删除 12
删除 14
删除 19
删除 17
删除 22
删除 20

测试结束