rbtree.cpp

红黑树的c++实现

#include<iostream>
#include<ctime>
#include<string>
using namespace std;
typedef int key_t;
typedef int data_t;
typedef enum color_t
{
	RED=0,
	BLACK=1
}color_t;
class rbtree
{
public:
	rbtree *left,*right,*parent;
	key_t key;
	data_t data;
	color_t color;
	rbtree* search(key_t key,rbtree *root);
	rbtree* insert(key_t key,data_t data,rbtree *root);
	rbtree* remove(key_t key,rbtree *root);
	rbtree* newnode(key_t key,data_t data);
	rbtree* rotate_left(rbtree *node,rbtree *root);
	rbtree* rotate_right(rbtree *node,rbtree *root);
	rbtree* insert_fixup(rbtree *node,rbtree *root);
	rbtree* remove_fixup(rbtree *node,rbtree *parent,rbtree *root);
	rbtree* search_auxiliary(key_t ket,rbtree *root,rbtree **save);

};
//-----------------------------------------------------------------------------------------------------------
int main()
{
    int i, count = 10000;
    key_t key;
	rbtree *root =NULL,*node =NULL;
	//srand(time(NULL));
	for(i=1;i<=count;i++)
	{
		key=rand();
		if((root=root->insert(key,i,root)))
		{
			cout<<"i= "<<i<<" insert "<<key<<" succeed!"<<endl;
		}
		else 
		{
			cout<<"i= "<<i<<" insert "<<key<<" failed!"<<endl;
			exit(-1);
		}
		if((node=root->search(key,root)))
		{
			cout<<"i= "<<i<<" search "<<key<<" succeed!"<<endl;
		}
		else
		{
			cout<<"i= "<<i<<" search "<<key<<" failed!"<<endl;
			exit(-1);
		}
		if(!(i%10))
		{
		    root=root->remove(key,root);
			cout<<"i= "<<i<<" remove "<<key<<" succeed!"<<endl;
		}
	}
	return 0;
}
//---------------------------------------------------------------------------------------------------------
rbtree* rbtree::newnode(key_t key,data_t data)
{
	rbtree *node = new rbtree;
	node->key=key;
	node->data=data;
	return node;
}
rbtree* rbtree::rotate_left(rbtree *node, rbtree *root)
{
	rbtree *right=node->right;
	if((node->right=right->left))
	{
		right->left->parent=node;
	}
	right->left=node;
	if ((right->parent = node->parent))
    {
        if (node == node->parent->right)
        {
            node->parent->right = right;
        }
        else
        {
            node->parent->left = right;
        }
    }
    else
    {
        root = right;
    }
    node->parent = right;

    return root;
}
rbtree* rbtree::rotate_right(rbtree *node,rbtree *root)
{
	rbtree* 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;
}
rbtree* rbtree::insert_fixup(rbtree *node, rbtree *root)
{
    rbtree *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 = rbtree::rotate_left(parent, root);
                    tmp = parent;
                    parent = node;
                    node = tmp;
                }

                parent->color = BLACK;
                gparent->color = RED;
				root = rbtree::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 = rbtree::rotate_right(parent, root);
                    tmp = parent;
                    parent = node;
                    node = tmp;
                }

                parent->color = BLACK;
                gparent->color = RED;
				root = rbtree::rotate_left(gparent, root);
            }
        }
    }

    root->color = BLACK;

    return root;
}
rbtree* rbtree::remove_fixup(rbtree *node, rbtree *parent, rbtree *root)
{
    rbtree *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 = rbtree::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 = rbtree::rotate_right(other, root);
                    other = parent->right;
                }
                other->color = parent->color;
                parent->color = BLACK;
                if (other->right)
                {
                    other->right->color = BLACK;
                }
				root = rbtree::rotate_left(parent, root);
                node = root;
                break;
            }
        }
        else
        {
            other = parent->left;
            if (other->color == RED)
            {
                other->color = BLACK;
                parent->color = RED;
				root = rbtree::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 = rbtree::rotate_left(other, root);
                    other = parent->left;
                }
                other->color = parent->color;
                parent->color = BLACK;
                if (other->left)
                {
                    other->left->color = BLACK;
                }
				root = rbtree::rotate_right(parent, root);
                node = root;
                break;
            }
        }
    }

    if (node)
    {
        node->color = BLACK;
    }

    return root;
}

rbtree* rbtree::search_auxiliary(key_t ket, rbtree *root, rbtree **save)
{
    rbtree *node = root, *parent = NULL;
    int ret;

    while (node)
    {
        parent = node;
        ret = node->key - key;
        if (0 < ret)
        {
            node = node->left;
        }
        else if (0 > ret)
        {
            node = node->right;
        }
        else
        {
            return node;
        }
    }

    if (save)
    {
        *save = parent;
    }

    return NULL;
}
rbtree* rbtree::insert(key_t key, data_t data, rbtree *root)
{
    rbtree *parent = NULL, *node;

    parent = NULL;
	if ((node = rbtree::search_auxiliary(key, root, &parent)))
    {
        return root;
    }
	node = rbtree::newnode(key, data);
    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 rbtree::insert_fixup(node, root);
}
rbtree* rbtree::search(key_t key, rbtree* root)
{
	return rbtree::search_auxiliary(key, root, NULL);
}

rbtree* rbtree::remove(key_t key,rbtree *root)
{
    rbtree *child, *parent, *old, *left, *node;
    color_t color;

	if (!(node = rbtree::search_auxiliary(key, root, NULL)))
    {
        cout<<"key does not exist!"<<endl;
        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 = rbtree::remove_fixup(child, parent, root);
    }

    return root;
}