红黑树源码（含注释）

Posted 2021-04-04 暗星涌动

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

typedef int key_t;

typedef enum color_t
{
    RED = 0,
    BLACK = 1
}color_t;

typedef struct rb_node_t
{
    struct rb_node_t *left, *right, *parent;
    key_t key;
    color_t color;
}rb_node_t;

rb_node_t* rb_insert(key_t key,  rb_node_t* root);
rb_node_t* rb_search(key_t key, rb_node_t* root);
rb_node_t* rb_erase(key_t key, rb_node_t* root);


// 打印 
int  Print(key_t e )    
{
    printf("%d ", e );
    return 1;
}

// 中序遍历 
int rb_print(rb_node_t* root,key_t(*Visit)(key_t))
{
  if(!root)
      return 0;
  else
  {
    if(root->left)
      rb_print(root->left,Visit);
    printf("%d ",root->key);
    if(root->right)
      rb_print(root->right,Visit);
  }
  return 1;
}


int main()
{
    int i, count = 10;
    key_t key;
    rb_node_t* root = NULL, *node = NULL;
 
// 随机插入元素    
    srand(time(NULL));
    for (i = 1; i < count; ++i)
    {
        key = rand() % count;  // key的值范围为 0 ~ 9 
        if ((root = rb_insert(key, root)))
        {
            printf("[i = %d] insert key %d success!\n", i, key);
        }
        else
        {
            printf("[i = %d] insert key %d error!\n", i, key);
            exit(-1);
        }
// 打印树的中序遍历序列 
        rb_print(root,Print);
        printf("\n");
    }

// 查找元素  
  key_t e;
  printf("The number e to be searched is: ");
  scanf("%d",&e); 
  if ((node = rb_search(e, root)))
        printf("Have searched e !\n");
    else
        printf("Can not find e!\n");

// 删除元素    
  printf("The number e to be erased is: ");
    scanf("%d",&e);
  rb_erase(e, root);
        
    return 0;
}


// 结点的初始化 
static rb_node_t* rb_new_node(key_t key)
{
    rb_node_t *node = (rb_node_t*)malloc(sizeof(struct rb_node_t));

    if (!node)  // 分配空间失败 
    {
        printf("malloc error!\n");  
        exit(-1);
    }
    node->key = key;

    return node;
}

// 左旋调节 
static rb_node_t* rb_rotate_left(rb_node_t* node, rb_node_t* root)
{
    rb_node_t* right = node->right;    // right 为 node 的右子结点 

    if ((node->right = right->left))  // right 的左子节点转给 node 当右子结点 
    {
        right->left->parent = node;    // 更改  right 的左子节点的父结点为 node 
    }
    right->left = node;    // node 变成 right 的左子节点 

    if ((right->parent = node->parent))    // right 的父结点改为 node 的父结点 
    {
        if (node == node->parent->right)  // 如果 node 是 父结点的右支，right 替代 node 
        {
            node->parent->right = right;  
        }
        else  // node 是父结点的左支 
        {
            node->parent->left = right;
        }
    }
    else  // node 是根节点，则 right 变为根节点 
    {
        root = right;
    }
    node->parent = right;  // node 的父结点改为其原右子结点 

    return root;
}

// 右旋调节 
static rb_node_t* rb_rotate_right(rb_node_t* node, rb_node_t* root)
{
    rb_node_t* left = node->left;  // 与左旋调节对称 

    if ((node->left = left->right))
    {
        left->right->parent = node;
    }
    left->right = node;

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

    return root;
}

// 插入元素后调节平衡 
static rb_node_t* rb_insert_rebalance(rb_node_t *node, rb_node_t *root)
{
    rb_node_t *parent, *gparent, *uncle, *tmp;

    while ((parent = node->parent) && parent->color == RED)   //父结点为红色
    {
        gparent = parent->parent;  // 祖父节点 

        if (parent == gparent->left)  // 父结点为祖父节点的左支 
        {
            uncle = gparent->right;  
            if (uncle && uncle->color == RED)  // 叔结点也为红色 
            {
                uncle->color = BLACK;  // 父结点和叔结点变为黑色，祖父节点变为红色 
                parent->color = BLACK;
                gparent->color = RED;
                node = gparent;
            }
            else  // 叔结点为黑色 
            {
                if (parent->right == node)  // 插入节点位于父结点的右支 
                {
                    root = rb_rotate_left(parent, root);  // 左旋 
                    tmp = parent;
                    parent = node;
                    node = tmp;    // 以父结点为支点再进行调整 
                }

                parent->color = BLACK;  // 父结点改为黑色，祖父节点改为红色 
                gparent->color = RED;
                root = rb_rotate_right(gparent, root);  // 右旋 
            }
        } 
        else  // 父结点为祖父节点的右支，调整与左支的对称 
        {
            uncle = gparent->left;
            if (uncle && uncle->color == RED)
            {
                uncle->color = BLACK;
                parent->color = BLACK;
                gparent->color = RED;
                node = gparent;
            }
            else
            {
                if (parent->left == node)
                {
                    root = rb_rotate_right(parent, root);
                    tmp = parent;
                    parent = node;
                    node = tmp;
                }

                parent->color = BLACK;
                gparent->color = RED;
                root = rb_rotate_left(gparent, root);
            }
        }
    }

    root->color = BLACK;

    return root;
}

//删除元素的平衡 
static rb_node_t* rb_erase_rebalance(rb_node_t *node, rb_node_t *parent, rb_node_t *root)
{
    rb_node_t *other, *o_left, *o_right;

    while ((!node || node->color == BLACK) && node != root)  // 当前节点是黑色 
    {
        if (parent->left == node)  // 当前节点位于父结点的左支 
        {
            other = parent->right;  
            if (other->color == RED)  // 兄弟节点是红色 
            {
                other->color = BLACK;  // 父结点和兄弟节点的颜色互换 
                parent->color = RED;
                root = rb_rotate_left(parent, root);  // 左旋 
                other = parent->right;
            }
            if ((!other->left || other->left->color == BLACK) &&(!other->right || other->right->color == BLACK)) // 兄弟节点的左右孩子为黑色 
            {
                other->color = RED;  // 兄弟节点变为红色，父结点再进行一次平衡操作 
                node = parent;
                parent = node->parent;
            }
            else
            {
                if (!other->right || other->right->color == BLACK)  // 只有兄弟节点的右孩子为黑色 
                {
                    if ((o_left = other->left))  // 左孩子存在，则为红色
                    {
                        o_left->color = BLACK;  // 左孩子变为黑色 
                    }
                    other->color = RED;  // 兄弟节点变为红色，以兄弟节点为支点右旋 
                    root = rb_rotate_right(other, root);
                    other = parent->right;
                }
                other->color = parent->color;  // 父结点的颜色赋给兄弟节点，父结点变为黑色 
                parent->color = BLACK;
                if (other->right)  // 兄弟节点的右孩子存在，则为红色 
                {
                    other->right->color = BLACK;  // 兄弟节点的右孩子变为黑色 
                }
                root = rb_rotate_left(parent, root);  // 左旋 
                node = root;
                break;
            }
        }
        else  // 当前节点位于父结点的右支，调整过程与左支对称 
        {
            other = parent->left;
            if (other->color == RED)
            {
                other->color = BLACK;
                parent->color = RED;
                root = rb_rotate_right(parent, root);
                other = parent->left;
            }
            if ((!other->left || other->left->color == BLACK) &&(!other->right || other->right->color == BLACK))
            {
                other->color = RED;
                node = parent;
                parent = node->parent;
            }
            else
            {
                if (!other->left || other->left->color == BLACK)
                {
                    if ((o_right = other->right))
                    {
                        o_right->color = BLACK;
                    }
                    other->color = RED;
                    root = rb_rotate_left(other, root);
                    other = parent->left;
                }
                other->color = parent->color;
                parent->color = BLACK;
                if (other->left)
                {
                    other->left->color = BLACK;
                }
                root = rb_rotate_right(parent, root);
                node = root;
                break;
            }
        }
    }

    if (node)
    {
        node->color = BLACK;
    } 

    return root;
}

//查找元素 
static rb_node_t* rb_search_auxiliary(key_t key, rb_node_t* root, rb_node_t** save)
{
    rb_node_t *node = root, *parent = NULL;
    int ret;

    while (node)
    {
        parent = node;
        ret = node->key - key;
        if (0 < ret)  // 结点的值比 key 值大，则从其左子树继续查找 
        {
            node = node->left;
        }
        else if (0 > ret)  // 结点的值比 key 值小，则从其右子树继续查找
        {
            node = node->right;
        }
        else  // 找到 key 值 
        {
            return node;
        }
    }

    if (save)
    {
        *save = parent;
    }

    return NULL;
}

//插入元素 
rb_node_t* rb_insert(key_t key, rb_node_t* root)
{
    rb_node_t *parent = NULL, *node;

    parent = NULL;
    if ((node = rb_search_auxiliary(key, root, &parent)))  // 树中已有该元素，则不再插入 
    {
        return root;
    }

    node = rb_new_node(key);  // 给结点赋值，并初始化其颜色为红色 
    node->parent = parent; 
    node->left = node->right = NULL;
    node->color = RED;

    if (parent)  // 选择插入到左支还是右支 
    {
        if (parent->key > key)
        {
            parent->left = node;
        }
        else
        {
            parent->right = node;
        }
    }
    else  // 该节点为根节点 
    {
        root = node;
    }

    return rb_insert_rebalance(node, root);
}

//查找元素 
rb_node_t* rb_search(key_t key, rb_node_t* root)
{
    return rb_search_auxiliary(key, root, NULL);
}

//删除元素 
rb_node_t* rb_erase(key_t key, rb_node_t *root)
{
    rb_node_t *child, *parent, *old, *left, *node;
    color_t color;

    if (!(node = rb_search_auxiliary(key, root, NULL)))
    {
        printf("e is not exist!\n");  // 没有找到该元素 
        return root;
    }

    old = node;

    if (node->left && node->right)  // 被删除结点左右子树存在 
    {
        node = node->right;
        while ((left = node->left) != NULL)  // 寻找当前节点，即被删除结点的中序后继结点 
        {
            node = left;
        }
        child = node->right;
        parent = node->parent;
        color = node->color;

        if (child)
        {
            child->parent = parent;
        }
        if (parent)
        {
            if (parent->left == node)  // 当前节点在左支 
            {
                parent->left = child; 
            }
            else  // 当前节点在右支
            {
                parent->right = child;
            }
        }
        else  
        {
            root = child;
        }

        if (node->parent == old)
        {
            parent = node;
        }

        node->parent = old->parent;   // 当前节点和根节点互换 
        node->color = old->color;
        node->right = old->right;
        node->left = old->left;

        if (old->parent)
        {
            if (old->parent->left == old)  // 被删除结点在左支 
            {
                old->parent->left = node;
            }
            else  // 被删除结点在右支
            {
                old->parent->right = node;
            }
        } 
        else  // 被删除结点为根节点 
        {
            root = node;
        }

        old->left->parent = node;   // 连接当前节点和被删除结点的子树 
        if (old->right)
        {
            old->right->parent = node;
        }
    }
    else
    {
        if (!node->left)  // 左子树不存在 
        {
            child = node->right;
        }
        else if (!node->right)  // 右子树不存在 
        {
            child = node->left;
        }
        parent = node->parent;
        color = node->color;

        if (child)
        {
            child->parent = parent;  // 替代被删除结点 
        }
        if (parent)
        {
            if (parent->left == node)  // 左支 
            {
                parent->left = child;
            }
            else  // 右支
            {
                parent->right = child;   
            }
        }
        else  // 根节点 
        {
            root = child;
        }
    }

    free(old);

    if (color == BLACK)
    {
        root = rb_erase_rebalance(child, parent, root);
    }
    
//打印删除后的中序遍历序列    
    printf("after deleting : ");
    rb_print(root,Print);
    printf("\n");

    return root;
}