红黑树源码（含注释）

小灰灰 2024-02-19 13:03 181阅读 0赞

## 红黑树源码（含注释） ##

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