rbtree.c

C/C++ 多任务下的数据结构与算法 （周伟明）华中科技大学出版社

/*
 * Copyright (c) 2006-2008
 * Author: Weiming Zhou
 *
 * Permission to use, copy, modify, distribute and sell this software
 * and its documentation for any purpose is hereby granted without fee,
 * provided that the above copyright notice appear in all copies and
 * that both that copyright notice and this permission notice appear
 * in supporting documentation.  
 */

#include <stdlib.h>
#include "CapiGlobal.h"
#include "BinTree.h"
#include "RBTree.h"

#pragma warning(disable : 4996)


/**	红黑树的创建函数

	@param	void - 无	
	@return	RBTREE * - 成功返回创建的红黑树指针，
                       失败返回NULL。
*/
RBTREE *RBTree_Create(void)
{
    RBTREE *pTree;

    pTree = (RBTREE *)malloc(sizeof(RBTREE));
    if ( pTree != NULL )
    {
        pTree->pRoot = NULL;
        pTree->uNodeCount = 0;
        pTree->pCursor = NULL;
    }

    return pTree;
}

/**	红黑树节点创建函数

	@param	void *pData - 节点数据指针	
	@return	static void * - 返回创建的节点指针	
*/
static void *RBTreeNode_Create(void *pData)
{
    RBTREENODE *pNewNode;

    pNewNode = (RBTREENODE *)malloc(sizeof(RBTREENODE));
    if ( pNewNode != NULL )
    {
        pNewNode->pLeft= NULL;
        pNewNode->pRight = NULL;
        pNewNode->nMagic = RBTREE_COLOR_RED;
        pNewNode->pData = pData;
    }
    return (void *)pNewNode;
}

/**	红黑树的释放函数，释放以指定节点为根节点的子树

	@param	RBTREENODE *pNode - 要释放的根节点	
	@param	DESTROYFUNC DestroyFunc - 数据释放回调函数	
	@return	static void - 无	
*/
static void RBTreeNode_Destroy(RBTREENODE *pNode, DESTROYFUNC DestroyFunc)
{
    if ( pNode != NULL )
    {
        if ( pNode->pLeft != NULL )
        {
            RBTreeNode_Destroy(pNode->pLeft, DestroyFunc);
        }
        if ( pNode->pRight != NULL )
        {
            RBTreeNode_Destroy(pNode->pRight, DestroyFunc);
        }
        if ( DestroyFunc != NULL && pNode->pData != NULL )
        {
            (*DestroyFunc)(pNode->pData);
        }
        free(pNode);
    }
}

/**	红黑树的释放函数

	@param	RBTREE *pTree - 红黑树指针	
	@param	DESTROYFUNC DestroyFunc -xxx	
	@return	void -xxxx	
*/
void RBTree_Destroy(RBTREE *pTree, DESTROYFUNC DestroyFunc)
{
    if ( pTree != NULL )
    {
        RBTreeNode_Destroy(pTree->pRoot, DestroyFunc);
        free( pTree );
    }
}

/**	红黑树的左旋操作，旋转前A节点和它的右节点B均存在
    旋转后A节点的右指针指向B节点的左节点，B节点的左指针指向A节点，

	@param	RBTREE *pTree -红黑树指针	
	@param	RBTREENODE *pANode -要旋转的节点	
	@return	static void -无	
*/
static void RBTree_RotateLeft(RBTREE *pTree, RBTREENODE *pANode)
{
    RBTREENODE  *pBNode; /* B节点指针 */

    pBNode = pANode->pRight;

    /* 将B节点的左节点变成A节点的右节点 */
    pANode->pRight = pBNode->pLeft;
    if ( pBNode->pLeft != NULL )
    {
        pBNode->pLeft->pParent = pANode;
    }

    /* 修改A节点的父指针和B节点的关系 */
    pBNode->pParent = pANode->pParent;
    if ( pANode == pTree->pRoot )
    {
        pTree->pRoot = pBNode;
    }
    else if ( pANode == pANode->pParent->pLeft )
    {
        pANode->pParent->pLeft = pBNode;
    }
    else
    {
        pANode->pParent->pRight = pBNode;
    }

    /* 将A节点变成B节点的左节点 */
    pBNode->pLeft = pANode;
    pANode->pParent = pBNode;
}


/**	红黑树的右旋操作，旋转前A节点和它的左节点B均存在
    旋转后A节点的左指针指向B节点的右节点，B节点的右指针指向A节点，

	@param	RBTREE *pTree -红黑树指针	
	@param	RBTREENODE *pANode -要旋转的节点	
	@return	static void -无	
*/
static void RBTree_RotateRight(RBTREE *pTree, RBTREENODE *pANode)
{
    RBTREENODE  *pBNode; /* B节点指针 */

    pBNode = pANode->pLeft;

    /* 将B节点的右节点变成A节点的左节点 */
    pANode->pLeft = pBNode->pRight;
    if ( pBNode->pRight != NULL )
    {
        pBNode->pRight->pParent = pANode;
    }

    /* 修改A节点的父指针和B节点的关系 */
    pBNode->pParent = pANode->pParent;
    if ( pANode == pTree->pRoot )
    {
        pTree->pRoot = pBNode;
    }
    else if ( pANode == pANode->pParent->pRight )
    {
        pANode->pParent->pRight = pBNode;
    }
    else
    {
        pANode->pParent->pLeft = pBNode;
    }

    /* 将A节点变成B节点的左节点 */
    pBNode->pRight = pANode;
    pANode->pParent = pBNode;
}

/**	红黑树的为插入操作调整颜色函数，

	@param	RBTREE *pTree - 红黑树指针	
	@param	RBTREENODE *pANode - 插入的节点指针	
	@return	static void - 无	
*/
static void RBTree_AdjustColorForInsert(RBTREE *pTree, RBTREENODE *pANode)
{
    /* 根据红黑树的算法，需要先将插入的A节点颜色预设成红色 */
    pANode->nMagic = RBTREE_COLOR_RED;

    while ( pANode != pTree->pRoot 
        && pANode->pParent->nMagic == RBTREE_COLOR_RED )
    {
        if ( pANode->pParent == pANode->pParent->pParent->pLeft )
        {
            /* 父节点为祖父节点左节点的情况 */

            RBTREENODE  *pNode = pANode->pParent->pParent->pRight;
            if ( pNode != NULL && pNode->nMagic == RBTREE_COLOR_RED )
            {
                /* LLr和LRr类型，直接修改颜色然后再继续循环调整祖父节点颜色 */
                pANode->pParent->nMagic = RBTREE_COLOR_BLACK;
                pNode->nMagic = RBTREE_COLOR_BLACK;
                pANode->pParent->pParent->nMagic = RBTREE_COLOR_RED;

                /* 由于祖父节点颜色被改变成了红色，因此祖父节点和它的父节点
                 * 有可能都是红色，这里需要继续循环来矫正祖父节点的颜色平衡情况 
                 */
                pANode = pANode->pParent->pParent;
            }
            else
            {
                /* LLb和LRb类型，直接通过旋转操作即可矫正颜色 */
                if ( pANode == pANode->pParent->pRight )
                {
                    pANode = pANode->pParent;
                    RBTree_RotateLeft(pTree, pANode);
                }
                pANode->pParent->nMagic = RBTREE_COLOR_BLACK;
                pANode->pParent->pParent->nMagic = RBTREE_COLOR_RED;

                RBTree_RotateRight(pTree, pANode->pParent->pParent);

                /* 这里由于A节点的父节点颜色是黑色的，实际上会退出循环
                 * 因此如果要优化代码大小的话, break 语句可以省略掉 
                 */
                break;
            }
        }
        else
        {
            /* 父节点为祖父节点右节点的情况 */

            RBTREENODE  *pNode = pANode->pParent->pParent->pLeft;
            if ( pNode != NULL && pNode->nMagic == RBTREE_COLOR_RED )
            {
                /* RLr和RRr类型，直接修改颜色然后再继续循环调整祖父节点颜色 */
                pANode->pParent->nMagic = RBTREE_COLOR_BLACK;
                pNode->nMagic = RBTREE_COLOR_BLACK;
                pANode->pParent->pParent->nMagic = RBTREE_COLOR_RED;

                /* 由于祖父节点颜色被改变成了红色，因此祖父节点和它的父节点
                 * 有可能都是红色，这里需要继续循环来矫正祖父节点的颜色平衡情况 
                 */
                pANode = pANode->pParent->pParent;
            }
            else
            {
                /* RLb和RRb类型，直接通过旋转操作即可矫正颜色 */
                if ( pANode == pANode->pParent->pLeft )
                {
                    pANode = pANode->pParent;
                    RBTree_RotateRight(pTree, pANode);
                }
                pANode->pParent->nMagic = RBTREE_COLOR_BLACK;
                pANode->pParent->pParent->nMagic = RBTREE_COLOR_RED;

                RBTree_RotateLeft(pTree, pANode->pParent->pParent);

                /* 这里由于A节点的父节点颜色是黑色的，实际上会退出循环
                 * 因此如果要优化代码大小的话, break 语句可以省略掉 
                 */
                break;
            }
        }
    }

    /* 在上面的颜色矫正过程中，根节点颜色可能会被改成红色 
     * 因此需要将根节点颜色重新赋为黑色即可
     */
    pTree->pRoot->nMagic = RBTREE_COLOR_BLACK;
}

/**	红黑树的插入函数

	@param	RBTREE *pTree - 红黑树指针	
	@param	void *pData - 要插入的数据指针	
	@param	COMPAREFUNC CompareFunc - 数据比较函数	
	@return	INT - CAPI_FAILED表示失败，CAPI_SUCCESS表示成功。	
*/
INT RBTree_Insert(RBTREE *pTree, void *pData, COMPAREFUNC CompareFunc)
{
    INT     nRet;
    RBTREENODE * pNode;
    
    nRet = CAPI_FAILED;

    pNode = RBTreeNode_Create(pData);
    if ( pNode != NULL )
    {
        nRet = RBTree_Inter_Insert(pTree, pNode, CompareFunc);
    }
    return nRet;
}

/**	红黑树的插入函数，主要为需要继承的模块而使用，
    一般的红黑树应用模块请不要调用这个函数

	@param	RBTREE *pTree - 红黑树指针
	@param	void *pData - 数据指针	
	@param	COMPAREFUNC CompareFunc - 比较函数	
	@param	void *pTreeNode - 插入的节点	
	@return	INT -CAPI_SUCCESS表示成功	
*/
INT RBTree_Inter_Insert(RBTREE *pTree,  RBTREENODE *pNewNode, 
                        COMPAREFUNC CompareFunc)
{
    RBTREENODE  *pNode;
    RBTREENODE  *pParentNode;
    INT         nRet = 0;

    pParentNode = NULL;
    pNode = pTree->pRoot;

    while ( pNode != NULL )
    {
        pParentNode = pNode;
        nRet = (*CompareFunc)(pNode->pData, pNewNode->pData);
        if ( nRet < 0 )
        {
            pNode = pNode->pRight;
        }
        else if ( nRet > 0 )
        {
            pNode = pNode->pLeft;
        }
        else
        {
            return CAPI_FAILED;
        }
    }

    if ( pParentNode == NULL )
    {
        /* 树为空的情况 */
        pTree->pRoot = pNewNode;
        pNewNode->pParent = NULL;
    }
    else 
    {
        if ( nRet < 0 )
        {
            pParentNode->pRight = pNewNode;
        }
        else 
        {
            pParentNode->pLeft = pNewNode;
        }
        pNewNode->pParent = pParentNode;
    }