首先展示测试结果
/*************************************************************************
> 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;
}