封装红黑树代码

原创已于 2022-11-24 14:33:43 修改 · 472 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#算法 #c++ #数据结构

于 2022-11-24 14:28:19 首次发布

二叉树专栏收录该内容

4 篇文章

订阅专栏

这篇博客介绍了C++中使用红黑树数据结构的实现，包括初始化、插入节点、删除节点、遍历以及大小计算等功能。作者展示了如何在遍历过程中安全删除节点，并提供了完整的C接口和回调函数。此外，还讨论了后续可能的封装方向，如转换为类和迭代器结构。

首先展示测试结果

/*************************************************************************
    > File Name: test_rbtree_api.cc
    > Author: hsz
    > Brief:
    > Created Time: Thu 24 Nov 2022 11:08:34 AM CST
 ************************************************************************/

#include <utils/rbtree_api.h>
#include <iostream>
#include <stdio.h>
#include <stdlib.h>

using namespace std;

struct RBData {
    rb_node rbNode;
    int data;
};

rb_root root;

#define rbdata_entry(node) rb_entry(node, RBData, rbNode)

void free_node(rb_node *node)
{
    RBData *data = rbdata_entry(node);
    printf("%s() free(%p) data = %d\n", __func__, data, data->data);
    delete data;
}

bool processData(rb_node *node)
{
    RBData *data = rbdata_entry(node);
    printf("\t%s() data = %d\n", __func__, data->data);
    if (data->data == 4) {
        free_node(rbtree_erase_node(&root, node)); // 测试在遍历时删除节点是否产生异常
    }
    return false;
}

int compare_key(rb_node *node, const void *seed)
{
    int key = *static_cast<const int *>(seed);
    RBData *data = rbdata_entry(node);
    if (data->data < key) {
        return -1;
    }
    if (data->data > key) {
        return 1;
    }

    return 0;
}

int compare_node(rb_node *left, rb_node *right)
{
    RBData *l = rbdata_entry(left);
    RBData *r = rbdata_entry(right);

    if (l->data < r->data) {
        return -1;
    }
    if (l->data > r->data) {
        return 1;
    }

    return 0;
}

int main(int argc, char **argv)
{
    rbtree_init(&root);
    rbtree_foreach(&root, processData);

    printf("\n--->rbtree size = %zu\n", rbtree_size(&root));

    printf("insert to(%p) 1, 3, 5\n", &root);
    RBData *newNode = new RBData;
    newNode->data = 1;
    rbtree_insert(&root, &newNode->rbNode, compare_node);

    newNode = new RBData;
    newNode->data = 3;
    rbtree_insert(&root, &newNode->rbNode, compare_node);

    newNode = new RBData;
    newNode->data = 5;
    rbtree_insert(&root, &newNode->rbNode, compare_node);

    printf("\nforeach rbtree(%p)\n", &root);
    rbtree_foreach(&root, processData);

    printf("\ninsert to(%p) 2, 4\n", &root);
    newNode = new RBData;
    newNode->data = 2;
    rbtree_insert(&root, &newNode->rbNode, compare_node);

    newNode = new RBData;
    newNode->data = 4;
    rbtree_insert(&root, &newNode->rbNode, compare_node);

    printf("\n--->rbtree size = %zu\n", rbtree_size(&root));

    printf("\nforeach rbtree(%p) and will erase data 4\n", &root);
    rbtree_foreach(&root, processData);

    printf("\n--->rbtree size = %zu\n", rbtree_size(&root));

    int num = 3;
    printf("\nerase %d from rbtree(%p)\n", num, &root);
    rbtree_erase(&root, &num, compare_key);

    printf("\nforeach rbtree(%p)\n", &root);
    rbtree_foreach(&root, processData);

    printf("\nclear rbtree\n");
    rbtree_clear(&root, free_node);
    printf("\n--->rbtree size = %zu\n", rbtree_size(&root));
    return 0;
}

在这里插入图片描述

目前是以C接口和回调形式给出，后续会封装成类和iterator，或将参考QMap的封装形式

完整代码

/*************************************************************************
    > File Name: rbtree_base.h
    > Author: hsz
    > Brief: linux rbtree source code
    > Created Time: Sat 11 Dec 2021 01:26:52 PM CST
 ************************************************************************/

#ifndef __RBTREE_H__
#define __RBTREE_H__

struct rb_node
{
    struct rb_node *rb_parent;
    struct rb_node *rb_right;
    struct rb_node *rb_left;
    char            rb_color;
#define RB_RED      0
#define RB_BLACK    1
};

struct rb_root
{
    struct rb_node *rb_node;
};

#define RB_ROOT (struct rb_root){ (struct rb_node *)0, }
#define rb_entry(ptr, type, member) \
    ((type *)((char *)(ptr)-(unsigned long)(&((type *)0)->member)))

#ifdef __cplusplus
extern "C"
{
#endif

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

/* Find logical next and previous nodes in a tree */
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 *);

/* Fast replacement of a single node without remove/rebalance/add/rebalance */
void rb_replace_node(struct rb_node *victim, struct rb_node *newnode,
                struct rb_root *root);

#ifdef __cplusplus
}
#endif

static inline void rb_link_node(struct rb_node *node, struct rb_node *parent,
                struct rb_node **link)
{
    node->rb_parent = parent;
    node->rb_color = RB_RED;
    node->rb_left = node->rb_right = (struct rb_node *)0;
    *link = node;
}

#endif // __RBTREE_H__

/*************************************************************************
    > File Name: rbtree_base.cpp
    > Author: hsz
    > Brief:
    > Created Time: Sun 12 Dec 2021 11:11:06 PM CST
 ************************************************************************/

#include "rbtree_base.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))
            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 (struct rb_node *)0;
    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 (struct rb_node *)0;
    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 *newnode,
             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 = newnode;
        else
            parent->rb_right = newnode;
    } else {
        root->rb_node = newnode;
    }
    if (victim->rb_left)
        victim->rb_left->rb_parent = newnode;
    if (victim->rb_right)
        victim->rb_right->rb_parent = newnode;

    /* Copy the pointers/colour from the victim to the replacement */
    *newnode = *victim;
}

/*************************************************************************
    > File Name: rbtree_api.h
    > Author: hsz
    > Brief:
    > Created Time: Wed 23 Nov 2022 11:10:12 AM CST
 ************************************************************************/

#ifndef __UTILS_RBTREE_API_H__
#define __UTILS_RBTREE_API_H__

#include "rbtree_base.h"
#include "utils.h"

EXTERN_C_BEGIN

void rbtree_init(rb_root *root);
void rbtree_clear(rb_root *root, void (*free_node_cb)(rb_node *));

/**
 * @brief 将node插入到红黑树root
 * 
 * @param root 红黑树
 * @param node 红黑树节点
 * @param compare_callback 比较回调。返回0表示相等，返回1往右子树放，返回-1往左子树放。
 * @return rb_node* 成功返回插入的节点地址，失败返回nullptr
 */
rb_node *rbtree_insert(rb_root *root, rb_node *newNode, int (*compare_callback)(rb_node *, rb_node *));

/**
 * @brief 从红黑树中删除等于seed的节点，但不释放此节点内存
 * 
 * @param root 红黑树根节点
 * @param seed 种子
 * @param equal_callback 判断相等回调
 * @return rb_node* 成功返回删除的节点，失败返回nullptr
 */
rb_node *rbtree_erase(rb_root *root, const void *seed, int (*compare_callback)(rb_node *, const void *));

rb_node *rbtree_erase_node(rb_root *root, rb_node *node);

/**
 * @brief 对红黑树进行有序遍历(在遍历中删除是不安全的?)
 * 
 * @param root 红黑树根节点
 * @param data_callback 数据处理回调. 从循环中退出返回true，不退出返回false
 * @return true 遍历成功
 * @return false 遍历失败
 */
bool rbtree_foreach(rb_root *root, bool (*data_callback)(rb_node *node));

/**
 * @brief 从红黑树中查找与seed相等的节点
 * 
 * @param root 红黑树根节点
 * @param seed 种子数据
 * @param compare_callback 同rbtree_erase的compare_callback
 * @return rb_node* 返回查找到的节点，失败返回nullptr
 */
rb_node *rbtree_search(rb_root *root, const void *seed, int (*compare_callback)(rb_node *, const void *));

size_t rbtree_size(rb_root *root);

EXTERN_C_END

#endif // __UTILS_RBTREE_API_H__

/*************************************************************************
    > File Name: rbtree_api.cpp
    > Author: hsz
    > Brief:
    > Created Time: Wed 23 Nov 2022 03:56:21 PM CST
 ************************************************************************/

#include "rbtree_api.h"
#include <assert.h>

void rbtree_init(rb_root *root)
{
    if (root) {
        root->rb_node = nullptr;
    }
}

void rbtree_clear(rb_root *root, void (*free_node_cb)(rb_node *))
{
    if (root == nullptr) {
        return;
    }
    struct rb_node *node = nullptr;
    while (node = root->rb_node) {
        rb_erase(node, root);
        if (eular_likely(free_node_cb != nullptr)) {
            free_node_cb(node);
        }
    }
}

rb_node *rbtree_insert(rb_root *root, rb_node *newNode, int (*compare_callback)(rb_node *, rb_node *))
{
    if (root == nullptr || compare_callback == nullptr) {
        return nullptr;
    }

    struct rb_node **node = &(root->rb_node);
    struct rb_node *parent = nullptr;
    bool exists = false;
    int compareResult = 0;
    while (nullptr != (*node)) {
        parent = *node;
        compareResult = compare_callback(parent, newNode);
        if (compareResult < 0) { // parent < newNode
            node = &(*node)->rb_right;
        } else if (compareResult > 0) {
            node = &(*node)->rb_left;
        } else {
            exists = true;
            break;
        }
    }

    if (!exists) {
        rb_link_node(newNode, parent, node);
        rb_insert_color(newNode, root);
    }

    return exists ? nullptr : newNode;
}

rb_node *rbtree_erase(rb_root *root, const void *seed, int (*compare_callback)(rb_node *, const void *))
{
    if (root == nullptr || seed == nullptr || compare_callback == nullptr) {
        return nullptr;
    }

    struct rb_node *eraseNode = rbtree_search(root, seed, compare_callback);
    if (eraseNode != nullptr) {
        rb_erase(eraseNode, root);
    }

    return eraseNode;
}

bool isValidNode(rb_root *root, rb_node *node)
{
    assert(root && node);
    rb_node *parent = node;
    bool valid = false;
    while (parent) {
        if (parent == root->rb_node) {
            valid = true;
            break;
        }
        if (parent == parent->rb_parent) { // 需确定根节点的rb_parent是否为空，不为空则会死循环
            break;
        }
        parent = parent->rb_parent;
    }

    return valid;
}

rb_node *rbtree_erase_node(rb_root *root, rb_node *node)
{
    if (root == nullptr || node == nullptr) {
        return nullptr;
    }
    assert(isValidNode(root, node));
    
    rb_erase(node, root);
    return node;
}

bool rbtree_foreach(rb_root *root, bool (*data_callback)(rb_node *node))
{
    if (root == nullptr || data_callback == nullptr) {
        return false;
    }

    struct rb_node *begin = rb_first(root);
    struct rb_node *next = begin;
    struct rb_node *end = rb_last(root);

    if (begin == nullptr) { // no nodes
        return true;
    }

    // NOTE 由于遍历时是有序的，故保存next后，即使在回调中删除begin节点也不会导致循环出现问题，因为下一个节点本就应该是next
    while (begin != end) {
        next = rb_next(next);
        if (data_callback(begin)) {
            return true;
        }
        begin = next;
    }

    data_callback(end);
    return true;
}

rb_node *rbtree_search(rb_root *root, const void *seed, int (*compare_callback)(rb_node *, const void *))
{
    if (root == nullptr || compare_callback == nullptr) {
        return nullptr;
    }

    struct rb_node *curr = root->rb_node;
    bool exist = false;
    int compareResult = 0;
    while (curr) {
        compareResult = compare_callback(curr, seed);
        if (compareResult < 0) {
            curr = curr->rb_left;
        } else if (compareResult > 0) {
            curr = curr->rb_right;
        } else {
            exist = true;
            break;
        }
    }

    return exist ? curr : nullptr;
}

size_t rbtree_size(rb_root *root)
{
    if (root == nullptr) {
        return 0;
    }

    size_t count = 0;
    struct rb_node *begin = rb_first(root);
    struct rb_node *next = begin;
    struct rb_node *end = rb_last(root);

    if (begin == nullptr) {
        return 0;
    }

    while (begin != end) {
        next = rb_next(next);
        begin = next;
        ++count;
    }

    ++count; // for end
    return count;
}