C语言版红黑树V2.0，从Linux内核中扒来的--速度巨快，支持select和rank

于 2011-12-01 21:54:31 发布
阅读量1.9k 收藏
点赞数
分类专栏： 算法 文章标签： linux内核 语言 c struct insert delete
 C语言版红黑树V2.0，从Linux内核中扒来的--速度巨快，支持select和rank
2011年12月01日 星期四 下午 09:04
转载自 分享
最终编辑 rangemq
把上一个版本改进了，现在支持size，select和rank了。
速度上比原来稍稍慢一点，但是比stl 的set还是快20倍左右，并且stl的set不支持select和rank。
代码：
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <map>
#include <set>
using namespace std;
int rand(int max) /*随机函数,用于生成一个0 - max-1的整数*/
{
    //srand((unsigned) time(0));
    return rand()%(max);
}


#define    RB_RED        0
#define    RB_BLACK    1
/* 定义自己的entry，first是key，second是value。*/
struct my_data{
    int first,second;
};
/*定义自己的数据的比较函数*/
static int cmp(struct my_data * left, struct my_data * right){
    return left->first - right->first;
}
struct rb_node {
    //在rb_node中添加自己的数据
    struct my_data key;
    struct rb_node *rb_parent;
    int rb_color;
    struct rb_node *rb_right;
    struct rb_node *rb_left;
    //used for select and rank function. when a new rb_node is initialized, weight should be 1.
    int weight;
};

struct rb_root {
    struct rb_node *rb_node;
    int size;
};

static void __rb_rotate_left(struct rb_node *node, struct rb_root *root)
{
    struct rb_node *right = node->rb_right;

    int a=0,b=0,c=0;
    if(node->rb_left)
        a=node->rb_left->weight;
    if(right->rb_left)
        b=right->rb_left->weight;
    if(right->rb_right)
        c=right->rb_right->weight;
    node->weight=a+b+1;
    right->weight=a+b+c+2;

    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;

    int a=0,b=0,c=0;
    if(node->rb_right)
        c=node->rb_right->weight;
    if(left->rb_left)
        a=left->rb_left->weight;
    if(left->rb_right)
        b=left->rb_right->weight;
    node->weight=b+c+1;
    left->weight=a+b+c+2;

    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)
{
    root->size--;
    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) != NULL)
            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;

        struct rb_node * q = parent;
        while(q){
            q->weight--;
            q=q->rb_parent;
        }

        node->rb_parent = old->rb_parent;
        node->rb_color = old->rb_color;
        node->weight = old->weight;
        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)
    {
        struct rb_node * q = parent;
        while(q){
            q->weight--;
            q=q->rb_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 NULL;
    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 NULL;
    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;
}

static int my_rb_insert(struct rb_node *t, struct rb_root *root)
{
    struct rb_node **p = &root->rb_node;
    struct rb_node *parent = NULL;
    int sig;

    while (*p) {
        parent = *p;
        parent->weight++;
        sig = cmp(&t->key, &parent->key);
        if (sig < 0)
            p = &(parent)->rb_left;
        else if (sig > 0)
            p = &(parent)->rb_right;
        else{
            while(parent){
                parent->weight--;
                parent=parent->rb_parent;
            }
            return 1;
        }
    }
    root->size++;
    t->rb_parent = parent;
    t->rb_color = RB_RED;
    t->rb_left = t->rb_right = NULL;
    t->weight=1;
    *p = t;
    rb_insert_color(t, root);
    return 0;
}

static struct rb_node * my_rb_find(struct my_data * key, struct rb_root *root)
{
    struct rb_node *n = root->rb_node;
    int sig;

    while (n) {
        sig = cmp(key, &n->key);
        if (sig < 0)
            n = n->rb_left;
        else if (sig > 0)
            n = n->rb_right;
        else
            return n;
    }
    return NULL;
}
/*order base is from 1 to root->size*/
static struct rb_node * select(int order, struct rb_root * root){
    if(order<=0 || order > root->size)
        return NULL;
    struct rb_node * cn=root->rb_node;
    int ls=0;
    while(cn){
        if(cn->rb_left)
            ls=cn->rb_left->weight;
        else ls=0;
        if(ls+1==order)
            return cn;
        else if(ls+1>order){
            cn=cn->rb_left;
        }else{
            order-=ls+1;
            cn=cn->rb_right;
        }
    }
    return NULL;
}
/*return a integer which is the number of rb_node->key <= w->key */
static int rank(struct my_data * w, struct rb_root * root){
    if(!root->rb_node)
        return 0;
    int res=0;
    struct rb_node * cn=root->rb_node;
    int sig;
    while(cn){
        sig = cmp(&cn->key,w);
        if(sig==0)
            return res+1+(cn->rb_left?cn->rb_left->weight:0);
        else if(sig<0){
            res+=1+(cn->rb_left?cn->rb_left->weight:0);
            cn = cn->rb_right;
        }else{
            cn = cn->rb_left;
        }
    }
    return res;
}
/*注意参数是rb_node * t,因此要删除一个元素，可能首先要my_rb_find。也可以在外部提供一个cache，
以省去my_rb_find的过程。
*/
static void my_rb_delete(struct rb_node *t, struct rb_root *root)
{
    rb_erase(t, root);
}
rb_node cache[1000000]; //for testing
int pos;
bool test_eq(set<int> & si, rb_root * root){
    rb_node * rbp = rb_first(root);
    set<int>::iterator it=si.begin();
    while(true){
        if(rbp == NULL && it == si.end())
            return true;
        if(rbp != NULL && it != si.end()){
            //printf("%d %d\n",rbp->key.first,*it);
            if(rbp->key.first != *it)
                return false;
            rbp = rb_next(rbp);
            it++;
        }else{
            return false;
        }
    }
    return true;
}
int getK(set<int> & si,int order){
    set<int>::iterator it=si.begin();
    for(int i=0;i<order-1;i++)
        it++;
    return *it;
}
int main(int argc, char **argv)
{
    //定义一个空的红黑树
    struct rb_root test1={0,};
    set<int> test2;
    int count=10000;
    int max_value=3000;
    pos=0;
    srand((unsigned) time(0));
    int * insert_value=new int[count];
    int * delete_value=new int[count];
    if(insert_value==NULL || delete_value ==NULL){
        printf(" not enough heap!\n");
        return 0;
    }
    for(int i=0;i<count;i++){
        insert_value[i]=rand(max_value);
        delete_value[i]=rand(max_value);
    }
    struct my_data tmp;
    clock_t t1=clock();
    for(int i=0;i<count ;i++){
        int next=insert_value[i];
        //printf("%d next:",next);
        cache[pos].key.first=next;
        //插入
        my_rb_insert(&cache[pos],&test1);
        pos++;
    }
    for(int j=0;j<count;j++){
        int next=delete_value[j];
        tmp.first=next;
        //查找
        rb_node * find = my_rb_find(&tmp,&test1);
        if(find){
            //删除
            my_rb_delete(find,&test1);
        }
    }
    clock_t t2=clock();
    printf("rb time: %ld\n",t2-t1);

    clock_t t3=clock();
    for(int i=0;i<count ;i++){
        int next=insert_value[i];
        test2.insert(next);
    }
    for(int j=0;j<count;j++){
        int next=delete_value[j];
        test2.erase(next);
    }
    clock_t t4=clock();
    printf("set time: %ld\n",t4-t3);


    printf(" set size: %d rb_size: %d\n",test2.size(),test1.size);
    if(test_eq(test2,&test1)==false){
        printf("error\n");
    }
    int times=1000;
    while(times-->0){
        int order=rand(test1.size)+1;
        rb_node * find =select(order, & test1);
        int k=getK(test2,order);
        int rankr=rank(&find->key,&test1);
        if(rankr!=order){
            printf("error!!!!!\n");
        }
        if(find->key.first!=k){
            printf("error: %d %d\n",find->key.first,k);
            break;
        }
    }

    system("pause");
    return 0;
}

wxl1986622
关注
  • 0
    点赞
  • 踩
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
专栏目录
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值