linux下红黑树使用实例

最近使用到红黑树写了一个电话本，使用的linux内核中的标准接口，然后加上了一些自己写的接口和测试程序，记录如下，方便后续使用，以下代码网上几乎都可以找到，只是根据自己所做的项目做了一点修正：

rbtree.h

/*  Red Black Trees  (C) 1999  Andrea Arcangeli <[e-mail]>    This program is free software; you can redistribute it and/or modify  it under the terms of the GNU General Public License as published by  the Free Software Foundation; either version 2 of the License, or  (at your option) any later version.  This program is distributed in the hope that it will be useful,  but WITHOUT ANY WARRANTY; without even the implied warranty of  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the  GNU General Public License for more details.  You should have received a copy of the GNU General Public License  along with this program; if not, write to the Free Software  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA  linux/include/linux/rbtree.h  To use rbtrees you'll have to implement your own insert and search cores.  This will avoid us to use callbacks and to drop drammatically performances.  I know it's not the cleaner way,  but in C (not in C++) to get  performances and genericity...  Some example of insert and search follows here. The search is a plain  normal search over an ordered tree. The insert instead must be implemented  in two steps: First, the code must insert the element in order as a red leaf  in the tree, and then the support library function rb_insert_color() must  be called. Such function will do the not trivial work to rebalance the  rbtree, if necessary. */#ifndef_LINUX_RBTREE_H#define_LINUX_RBTREE_H#define MAX_LEN 32  #define RECORD_MAX 512//定义数据记录  struct PBookData  {     char name[MAX_LEN];     char num1[MAX_LEN];    char num2[MAX_LEN]; }RecordData[RECORD_MAX];typedef struct PBookData  st_PBookRecord;  struct rb_node{unsigned long  rb_parent_color;#defineRB_RED0#defineRB_BLACK1struct rb_node *rb_right;struct rb_node *rb_left;} __attribute__((aligned(sizeof(long))));struct rb_root{struct rb_node *rb_node;};//自定义类型struct mytype {      struct rb_node my_node;      //int num;  st_PBookRecord *my_pbook;};  typedef struct mytype *  ElemDataNode;//取r节点的父节点的地址#define rb_parent(r)   ((struct rb_node *)((r)->rb_parent_color & ~3)) //获取r节点的颜色：1是黑，0是红。#define rb_color(r)   ((r)->rb_parent_color & 1)//测试节点是否为红色#define rb_is_red(r)   (!rb_color(r))//测试节点是否为黑色#define rb_is_black(r) rb_color(r)#define rb_set_red(r)  do { (r)->rb_parent_color &= ~1; } while (0)#define rb_set_black(r)  do { (r)->rb_parent_color |= 1; } while (0)#define offsetof(TYPE, MEMBER) ((size_t) &((TYPE *)0)->MEMBER)#define container_of(ptr, type, member) ({                \    const typeof( ((type *)0)->member ) *__mptr = (ptr);  \    (type *)( (char *)__mptr - offsetof(type,member) );})//设置节点父节点地址为p所指地址static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p){rb->rb_parent_color = (rb->rb_parent_color & 3) | (unsigned long)p;}//设置节点colorstatic inline void rb_set_color(struct rb_node *rb, int color){rb->rb_parent_color = (rb->rb_parent_color & ~1) | color;}#define RB_ROOT(struct rb_root) { NULL, }//初始根节点指针#definerb_entry(ptr, type, member) container_of(ptr, type, member)//包含ptr的结构体指针#define RB_EMPTY_ROOT(root)((root)->rb_node == NULL)//判断树是否空#define RB_EMPTY_NODE(node)(rb_parent(node) == node)//判断节点是否空，父亲是否等于自身#define RB_CLEAR_NODE(node)(rb_set_parent(node, node))//设置节点为空，父亲等于自身int rb_pbook_cmp(const char *s1, const char *s2);void rb_pbook_cp(st_PBookRecord * src, st_PBookRecord * dest);struct mytype *rb_search(struct rb_root *root, const char *name);int rb_insert(struct rb_root *root, struct mytype *data);void rb_delete(struct rb_root *root, const char *name);void rb_print(struct rb_root *tree);void rb_free(struct rb_root *tree);//把新插入的节点进行着色，并且修正红黑树使其达到平衡，其效果就是前文的insertFixup的效果void rb_insert_color(struct rb_node *, struct rb_root *);//To remove an existing node from a treevoid rb_erase(struct rb_node *, struct rb_root *);/*rb_next:返回node在书中的后继，这个稍微复杂一点。如果node的右孩子不为空，他只要返回node的右子树中最小的节点即可；如果为空，他要向上查找，找到迭代节点是其父亲的左孩子的节点，返回父节点。如果一直上述到了根节点，返回NULL。*/struct rb_node *rb_pgdn(struct rb_root *tree, const struct rb_node *node);struct rb_node *rb_pgup(struct rb_root *tree, const struct rb_node *node);struct rb_node *rb_next(const struct rb_node *);struct rb_node *rb_prev(const struct rb_node *);struct rb_node *rb_first(const struct rb_root *);struct rb_node *rb_last(const struct rb_root *);/* Fast replacement of a single node without remove/rebalance/add/rebalance *///To replace an existing node in a tree with a new one with the same key//void rb_replace_node(struct rb_node *victim, struct rb_node *new, struct rb_root *root);int rb_replace_node(const char *name, struct rb_node *new, struct rb_root *root);//插入节点时需要把新节点指向其父亲节点static inline void rb_link_node(struct rb_node * node, struct rb_node * parent, struct rb_node ** rb_link){node->rb_parent_color = (unsigned long )parent;node->rb_left = node->rb_right = NULL;*rb_link = node;}#endif/* _LINUX_RBTREE_H */

rbtree.c

/*  Red Black Trees  (C) 1999  Andrea Arcangeli <[e-mail]>  (C) 2002  David Woodhouse <[e-mail]>    This program is free software; you can redistribute it and/or modify  it under the terms of the GNU General Public License as published by  the Free Software Foundation; either version 2 of the License, or  (at your option) any later version.  This program is distributed in the hope that it will be useful,  but WITHOUT ANY WARRANTY; without even the implied warranty of  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the  GNU General Public License for more details.  You should have received a copy of the GNU General Public License  along with this program; if not, write to the Free Software  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA  linux/lib/rbtree.c*/#include <stdio.h>#include <stdlib.h>#include "DataTree.h"#define DATA_FILE  "/etc_asterisk/pb/db.dat"  int rb_pbook_cmp(const char *s1, const char *s2){return strcmp(s1, s2);}void rb_pbook_cp(st_PBookRecord * src, st_PBookRecord * dest){strcpy(src->name, dest->name);strcpy(src->num1, dest->num1);strcpy(src->num2, dest->num2);}struct mytype *rb_search(struct rb_root *root, const char *name)  {      struct rb_node *node = root->rb_node;        while (node) {  struct mytype *data = container_of(node, struct mytype, my_node);     if (rb_pbook_cmp(name, data->my_pbook->name) < 0)node = node->rb_left;   else if (rb_pbook_cmp(name, data->my_pbook->name) > 0)node = node->rb_right;  else  return data;      }            return NULL;  }    int rb_insert(struct rb_root *root, struct mytype *data)  {  /*declare tmp rb_node ** in order to store parent address*/    struct rb_node **tmp = &(root->rb_node), *parent = NULL;        /* Figure out where to put new node */      while (*tmp) {  struct mytype *this = container_of(*tmp, struct mytype, my_node);    parent = *tmp;  if (rb_pbook_cmp(data->my_pbook->name, this->my_pbook->name) < 0) tmp = &((*tmp)->rb_left);  else if (rb_pbook_cmp(data->my_pbook->name, this->my_pbook->name) > 0)tmp = &((*tmp)->rb_right);  else   return -1;      }            /* Add new node and rebalance tree. */      rb_link_node(&data->my_node, parent, tmp);      rb_insert_color(&data->my_node, root);            return 0;  }    void rb_delete(struct rb_root *root, const char *name)  {       struct mytype *data = rb_search(root, name);      if (!data) {   fprintf(stderr, "Not found %s.\n", name);  return;      }            rb_erase(&(data->my_node), root);  free(data->my_pbook);    free(data);  }   void rb_print(struct rb_root *tree)  {      struct rb_node *node;            for (node = rb_first(tree); node; node = rb_next(node)) {struct mytype * data = rb_entry(node, struct mytype, my_node);printf("%s,%s,%s ", data->my_pbook->name, data->my_pbook->num1, data->my_pbook->num2); printf("0x%x,%s\n", node->rb_parent_color, rb_color(node)? "black":"red");}          printf("\n");  }  void rb_free(struct rb_root *tree)  {      struct rb_node *node;            for (node = rb_first(tree); node; node = rb_next(node)) {struct mytype * data = rb_entry(node, struct mytype, my_node);free(data->my_pbook);rb_erase(&(data->my_node), tree);free(data); }          printf("\n");  } static void __rb_rotate_left(struct rb_node *node, struct rb_root *root){struct rb_node *right = node->rb_right;  struct rb_node *parent = rb_parent(node);if ((node->rb_right = right->rb_left))rb_set_parent(right->rb_left, node);right->rb_left = node;rb_set_parent(right, parent);if (parent){if (node == parent->rb_left)parent->rb_left = right;elseparent->rb_right = right;}elseroot->rb_node = right;rb_set_parent(node, right);}static void __rb_rotate_right(struct rb_node *node, struct rb_root *root){struct rb_node *left = node->rb_left;struct rb_node *parent = rb_parent(node);if ((node->rb_left = left->rb_right))rb_set_parent(left->rb_right, node);left->rb_right = node;rb_set_parent(left, parent);if (parent){if (node == parent->rb_right)parent->rb_right = left;elseparent->rb_left = left;}elseroot->rb_node = left;rb_set_parent(node, left);}void rb_insert_color(struct rb_node *node, struct rb_root *root){struct rb_node *parent, *gparent;while ((parent = rb_parent(node)) && rb_is_red(parent)){gparent = rb_parent(parent);if (parent == gparent->rb_left){{register struct rb_node *uncle = gparent->rb_right;if (uncle && rb_is_red(uncle)){rb_set_black(uncle);rb_set_black(parent);rb_set_red(gparent);node = gparent;continue;}}if (parent->rb_right == node){register struct rb_node *tmp;__rb_rotate_left(parent, root);tmp = parent;parent = node;node = tmp;}rb_set_black(parent);rb_set_red(gparent);__rb_rotate_right(gparent, root);} else {{register struct rb_node *uncle = gparent->rb_left;if (uncle && rb_is_red(uncle)){rb_set_black(uncle);rb_set_black(parent);rb_set_red(gparent);node = gparent;continue;}}if (parent->rb_left == node){register struct rb_node *tmp;__rb_rotate_right(parent, root);tmp = parent;parent = node;node = tmp;}rb_set_black(parent);rb_set_red(gparent);__rb_rotate_left(gparent, root);}}rb_set_black(root->rb_node);}static void __rb_erase_color(struct rb_node *node, struct rb_node *parent, struct rb_root *root){struct rb_node *other;while ((!node || rb_is_black(node)) && node != root->rb_node){if (parent->rb_left == node){other = parent->rb_right;if (rb_is_red(other)){rb_set_black(other);rb_set_red(parent);__rb_rotate_left(parent, root);other = parent->rb_right;}if ((!other->rb_left || rb_is_black(other->rb_left)) &&    (!other->rb_right || rb_is_black(other->rb_right))){rb_set_red(other);node = parent;parent = rb_parent(node);}else{if (!other->rb_right || rb_is_black(other->rb_right)){rb_set_black(other->rb_left);rb_set_red(other);__rb_rotate_right(other, root);other = parent->rb_right;}rb_set_color(other, rb_color(parent));rb_set_black(parent);rb_set_black(other->rb_right);__rb_rotate_left(parent, root);node = root->rb_node;break;}}else{other = parent->rb_left;if (rb_is_red(other)){rb_set_black(other);rb_set_red(parent);__rb_rotate_right(parent, root);other = parent->rb_left;}if ((!other->rb_left || rb_is_black(other->rb_left)) &&    (!other->rb_right || rb_is_black(other->rb_right))){rb_set_red(other);node = parent;parent = rb_parent(node);}else{if (!other->rb_left || rb_is_black(other->rb_left)){rb_set_black(other->rb_right);rb_set_red(other);__rb_rotate_left(other, root);other = parent->rb_left;}rb_set_color(other, rb_color(parent));rb_set_black(parent);rb_set_black(other->rb_left);__rb_rotate_right(parent, root);node = root->rb_node;break;}}}if (node)rb_set_black(node);}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的最小节点，记为node。node = node->rb_right;while ((left = node->rb_left) != NULL)node = left;//node为后继if (rb_parent(old)) {if (rb_parent(old)->rb_left == old)rb_parent(old)->rb_left = node;elserb_parent(old)->rb_right = node;} elseroot->rb_node = node;child = node->rb_right;//child可能为NULLparent = rb_parent(node);//入参node的后继的父亲记为parent，parent此时是入参node的父亲的地址。color = rb_color(node);     //删除的颜色if (parent == old) {//y是z的right的情况。parent = node;} else {if (child)rb_set_parent(child, parent);parent->rb_left = child;node->rb_right = old->rb_right;       rb_set_parent(old->rb_right, node);}node->rb_parent_color = old->rb_parent_color;   //要保留的颜色。node->rb_left = old->rb_left;rb_set_parent(old->rb_left, node);goto color;}parent = rb_parent(node);color = rb_color(node);if (child)rb_set_parent(child, parent);if (parent){if (parent->rb_left == node)parent->rb_left = child;elseparent->rb_right = child;}elseroot->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(const struct rb_root *root){struct rb_node*n;n = root->rb_node;if (!n)return NULL;while (n->rb_left)n = n->rb_left;return n;}struct rb_node *rb_last(const struct rb_root *root){struct rb_node*n;n = root->rb_node;if (!n)return NULL;while (n->rb_right)n = n->rb_right;return n;}struct rb_node *rb_next(const struct rb_node *node){struct rb_node *parent;if (rb_parent(node) == node)return NULL;/* 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 (struct rb_node *)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 ((parent = rb_parent(node)) && node == parent->rb_right)node = parent;return parent;}struct rb_node *rb_prev(const struct rb_node *node){struct rb_node *parent;if (rb_parent(node) == node)return NULL;/* 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 (struct rb_node *)node;}/* No left-hand children. Go up till we find an ancestor which is a right-hand child of its parent */while ((parent = rb_parent(node)) && node == parent->rb_left)node = parent;return parent;}struct rb_node *rb_pgdn(struct rb_root *tree, const struct rb_node *node){struct rb_node *parent;/*determine whether the node is empty*/if (rb_parent(node) == node)return NULL;if (node == rb_last(tree)){node = rb_first(tree);return (struct rb_node *)node;}if (node->rb_right) {node = node->rb_right; while (node->rb_left)node=node->rb_left;return (struct rb_node *)node;}/*the scene node is in left child branch and biggest*/while ((parent = rb_parent(node)) && node == parent->rb_right)node = parent;return parent;}struct rb_node *rb_pgup(struct rb_root *tree, const struct rb_node *node){struct rb_node *parent;if (rb_parent(node) == node)return NULL;if (node == rb_first(tree)){node = rb_last(tree);return (struct rb_node *)node;}if (node->rb_left) {node = node->rb_left; while (node->rb_right)node=node->rb_right;return (struct rb_node *)node;}while ((parent = rb_parent(node)) && node == parent->rb_left)node = parent;return parent;}int rb_replace_node(const char *name, struct rb_node *new, struct rb_root *root){struct mytype *data = rb_search(root, name);if(data == NULL){printf("There are not exist %s.\n", name); return -1;} struct rb_node *victim = &data->my_node; struct rb_node *parent = rb_parent(victim);/* Set the surrounding nodes to point to the replacement */if (parent) {if (victim == parent->rb_left)parent->rb_left = new;elseparent->rb_right = new;} else {root->rb_node = new;}if (victim->rb_left)rb_set_parent(victim->rb_left, new);if (victim->rb_right)rb_set_parent(victim->rb_right, new);/* Copy the pointers/colour from the victim to the replacement */*new = *victim;return 0;}ElemDataNode rb_create_node(const char *name, const char *num1, const char *num2)  {  struct mytype *tmp;tmp = (struct mytype *)malloc(sizeof(struct mytype));  if (!tmp){perror("Allocate dynamic memory(mytype)");return NULL; }tmp->my_pbook = (struct PBookData *)malloc(sizeof(struct PBookData));if(tmp->my_pbook == NULL){perror("Allocate dynamic memory(st_PBookRecord)");return NULL; }   memset(tmp->my_pbook,0,sizeof(struct PBookData));strcpy(tmp->my_pbook->name, name);strcpy(tmp->my_pbook->num1, num1);strcpy(tmp->my_pbook->num2, num2);return tmp;  } int rb_write_data(struct rb_root *tree)  {  struct rb_node *node; st_PBookRecord  record;    FILE *pFile;   pFile = fopen(DATA_FILE,"w+");      if(pFile == NULL)      {          perror("Create File fail");          return 0;      }   //fseek(pFile, 0, SEEK_SET);    for (node = rb_first(tree); node; node = rb_next(node)) {struct mytype * data = rb_entry(node, struct mytype, my_node);if(fwrite(data->my_pbook,sizeof(struct PBookData),1,pFile) != 1) perror("file write fail");}    fclose(pFile) ;        return 0;} int rb_read_data(struct rb_root *tree)  {      FILE *pFile;       int i;  struct PBookData  record;    pFile = fopen(DATA_FILE,"r+");      if(pFile == NULL)      {          perror("/etc/db.dat");          return 0;      }      for(i=0; i<RECORD_MAX; i++)  {   if(fread(&record,sizeof(struct PBookData),1,pFile) == 0)break; rb_insert(tree, rb_create_node(record.name, record.num1, record.num2)); }     fclose(pFile) ;        return 0;} int main(int argc, char *argv[])  {  struct rb_root mytree = RB_ROOT; int ret;rb_read_data(&mytree);#if 0    printf("Test start:\n");   rb_insert(&mytree, rb_create_node("jix", "123", "456"));  rb_insert(&mytree, rb_create_node("hai", "234", "456"));  rb_insert(&mytree, rb_create_node("ho", "345", "456"));  rb_insert(&mytree, rb_create_node("liy", "456", "456"));  rb_insert(&mytree, rb_create_node("qiy", "567", "456"));  rb_insert(&mytree, rb_create_node("yaj", "789", "456"));  #endif    printf("\nthe first test\n");      rb_print(&mytree);printf("\nthe second test\n");     rb_delete(&mytree, "qiying"); //rb_insert(&mytree, rb_create_node("yajie", "789", "456")); rb_print(&mytree);ret = rb_replace_node("yajie", &(rb_create_node("qiying", "567", "456")->my_node), &mytree);if(ret < 0){goto ClearSpace;} rb_print(&mytree);rb_write_data(&mytree);ClearSpace:rb_free(&mytree);if(!RB_EMPTY_ROOT(&mytree))mytree.rb_node = NULL;    return 0;  }