红黑树的基本操作

#ifndef RBTREE_H_INCLUDED
#define RBTREE_H_INCLUDED

#undef NULL
#if defined(__cplusplus)
    #define NULL 0
#else
    #define NULL ((void *)0)
#endif

/*
红黑树是二叉查找树的一种且具有以下性质：
1：每个节点要么是红色要么是黑色
2：根节点和叶子节点都是黑色的
3：如果一个结点是红的，那么它的两个儿子都是黑的
4：每一条路径上黑节点的数目都相同
*/

/*
左旋：
right是node的右孩子。-----条件
旋转之后node成为right的左孩子。
right的左孩子成为node的右孩子。
node的左孩子不变，right的右孩子不变。
   node           right
   / \    ==>     / \
   a  right     node  y
       / \       / \
       b  y     a   b

*/

/*
右旋：
       node            left
       / \             / \
    left  y   ==>    a    node
   / \                    / \
  a   b                  b   y
父节点的左孩子才能右旋
右旋后父节点和左子树关系交换。
父节点变成左孩子的右孩子
左孩子的左孩子（a）位置不变
父节点的右孩子（y）位置不变
*/

typedef enum en_color
{
    RED = 0,
    BLK
}COLOR;


typedef struct tag_rb_t
{
    struct tag_rb_t *pstLt;     //左孩子节点
    struct tag_rb_t *pstRt;     //右孩子节点
    struct tag_rb_t *pstPt;     //双亲节点
    COLOR color;                //节点的颜色
    int key;                    //key
}rb_t;

//左旋
void rb_LeftRotate(rb_t* pNode, rb_t** ppRoot);
//右旋
void rb_RightRotate(rb_t* pNode, rb_t** ppRoot);
//插入时修正左子树
void rb_InsertFixupLeft(rb_t *pGrand, rb_t **pParent, rb_t *pUncle, rb_t **ppNode, rb_t **ppRoot);
//插入时修正右子树
void rb_InsertFixupRight(rb_t *pGrand, rb_t **pParent, rb_t *pUncle, rb_t **ppNode, rb_t **ppRoot); 
//插入时修正左右子树的入口
void rb_InsertFixup(rb_t *pNode, rb_t **ppRoot);
//插入
int rb_Insert(rb_t *pNode, rb_t **ppRoot);
//删除节点
void rb_Delete(rb_t *pDel, rb_t **ppRoot);
//查找一个节点
rb_t* rb_Search(int key);

#endif // RBTREE_H_INCLUDED

#include "stdafx.h"
#include "rbTree.h"

#define RETURN_ERR -1
#define RETURN_OK 0

/*****************************************************************************
函数功能：红黑树插入节点时左旋
函数入参：rb_t* pNode   左旋操作前的父节点
函数出参：rb_t** ppRoot 红黑树的根节点，在左旋过程中有可能需要改变根节点的位置
函数返回值：无
其他：左旋操作参看头文件中的说明
******************************************************************************/
void rb_LeftRotate(rb_t* pNode, rb_t** ppRoot)
{
    //pNode的右孩子。左旋发生在父节点（pNode）和右孩子（ppstRt）之间。左旋完成之后ppstRt成为父节点，pNode成为左孩子
    rb_t *ppstRt = pNode->pstRt;    

    //ppstRt的左孩子需要成为pNode的右孩子，pNode的左孩子不变，ppstRt的右孩子不变。
    if (NULL != (pNode->pstRt = ppstRt->pstLt))
    {
        ppstRt->pstLt->pstPt = pNode;   //ppstRt的左孩子成为node的右孩子，所以ppstRt的左孩子的父节点需要由ppstRt修改成pNode
    }
    ppstRt->pstLt = pNode;  //交换pNode和ppstRt的父子关系。


    if (NULL != (ppstRt->pstPt = pNode->pstPt))     //修改父节点。如果不为空说明pNode不是红黑树的根节点
    {
        //左旋后，未变动的节点的父节点也需要修改
        if (pNode == pNode->pstPt->pstRt)
        {
            pNode->pstPt->pstRt = ppstRt;   
        }
        else
        {
            pNode->pstPt->pstLt = ppstRt;
        }
    }
    else
    {
        *ppRoot = ppstRt;       //pNode是红黑树的根节点，此时根节点需要修改成为左旋前pNode的右孩子
    }
    pNode->pstPt = ppstRt;

    return;
}

/*****************************************************************************
函数功能：红黑树插入节点时右旋
函数入参：rb_t* pNode    右旋前的父节点
函数出参：rb_t** ppRoot  红黑树的根节点
函数返回值：无
其他：右旋操作参看头文件中的说明
******************************************************************************/
void rb_RightRotate(rb_t* pNode, rb_t** ppRoot)
{
    rb_t *ppstLt = pNode->pstLt;

    if (NULL != (pNode->pstLt = ppstLt->pstRt))
    {
        ppstLt->pstRt->pstRt = pNode;
    }
    ppstLt->pstRt = pNode;

    if (NULL != (ppstLt->pstPt = pNode->pstPt))
    {
        if (pNode == pNode->pstPt->pstRt)
        {
            pNode->pstPt->pstRt = ppstLt;
        } 
        else
        {
            pNode->pstPt->pstLt = ppstLt;
        }
    } 
    else
    {
        *ppRoot = ppstLt;
    }
    pNode->pstPt = ppstLt;

    return;
}

/*****************************************************************************
函数功能：红黑树插入节点时调整左子树
函数入参：rb_t *pGrand,  祖父节点
         rb_t **ppstPt, 父节点
         rb_t* *pUncle,  父节点的兄弟节点
         rb_t **ppNode,  待调整的节点
         rb_t **ppRoot   红黑树的根节点
函数出参：无
函数返回值：无
其他：
******************************************************************************/
void rb_InsertFixupLeft(rb_t *pGrand, rb_t **pppstPt, rb_t *pUncle, rb_t **ppNode, rb_t **ppRoot)
{
    rb_t *pNodeTmp = NULL;
    bool IsTrue = false;

    IsTrue = ((NULL != pUncle) && (RED == pUncle->color));      //uncle存在且是红色
    if (IsTrue)
    {
        pUncle->color = BLK;
        (*pppstPt)->color = BLK;
        pGrand->color = RED;
        (*ppNode) = pGrand;     //将祖父当做新增结点z，指针z上移俩层，且着为红色。
        return;
    }

    //uncle不存在，或者是黑色的
    if ((*ppNode) == (*pppstPt)->pstRt)    //pNode是右孩子，左旋的条件
    {
        rb_LeftRotate(*pppstPt, ppRoot);
        pNodeTmp = *pppstPt;
        *pppstPt = *ppNode;
        *ppNode = pNodeTmp;
    }

    //uncle是黑色的，此时pNode成为了左孩子。  
    (*pppstPt)->color = BLK;
    pGrand->color = RED;
    rb_RightRotate(pGrand, ppRoot);

    return;
}

/*****************************************************************************
函数功能：红黑树插入节点时调整右子树
函数入参：rb_t *pGrand,  祖父节点
         rb_t **ppstPt, 父节点
         rb_t* *pUncle,  父节点的兄弟节点
         rb_t **ppNode,  待调整的节点
         rb_t **ppRoot   红黑树的根节点
函数出参：无
函数返回值：无
其他：
******************************************************************************/
void rb_InsertFixupRight(rb_t *pGrand, rb_t **pppstPt, rb_t *pUncle, rb_t **ppNode, rb_t **ppRoot)
{
    rb_t *pNodeTmp = NULL;
    bool IsTrue;

    IsTrue = ((NULL != pUncle) && (RED == pUncle->color));
    if (IsTrue)
    {
        pUncle->color = BLK;
        (*pppstPt)->color = BLK;
        pGrand->color = RED;
        (*ppNode) = pGrand;
        return;
    }

    if ((*ppNode) == (*pppstPt)->pstLt)
    {
        rb_RightRotate(*pppstPt, ppRoot);
        pNodeTmp = *pppstPt;
        *pppstPt = *ppNode;
        *ppNode = pNodeTmp;
    }

    (*pppstPt)->color = BLK;
    pGrand->color = RED;
    rb_LeftRotate(pGrand, ppRoot);

    return;
}

/*****************************************************************************
函数功能：红黑树插入节点时调整子树
函数入参：rb_t *pNode,   带插入节点
         rb_t **ppRoot  红黑树的根节点
函数出参：无
函数返回值：无
其他：
******************************************************************************/
void rb_InsertFixup(rb_t *pNode, rb_t **ppRoot)
{
    rb_t *ppstPt = NULL;
    rb_t *pGrand = NULL;
    rb_t *pUncle = NULL;

    while ((NULL != (ppstPt = pNode->pstPt)) && (RED == ppstPt->color))
    {
        //ppstPt是pNode的父节点，且父节点是红色的
        pGrand = ppstPt->pstPt;
        if (ppstPt == pGrand->pstLt)
        {
            pUncle = pGrand->pstRt;
            rb_InsertFixupLeft(pGrand, &ppstPt, pUncle, &pNode, ppRoot);
        }
        else
        {
            pUncle = pGrand->pstLt;
            rb_InsertFixupRight(pGrand, &ppstPt, pUncle, &pNode, ppRoot);
        }
    }

    (*ppRoot)->color = BLK;

    return;
}

/*****************************************************************************
函数功能：红黑树插入节点
函数入参：rb_t *pNode,   插入节点
         rb_t **ppRoot  红黑树的根节点
函数出参：无
函数返回值：无
其他：
******************************************************************************/
int rb_Insert(rb_t *pNode, rb_t **ppRoot)
{
    rb_t **ppNodeTmp = ppRoot;
    rb_t *ppstPt = NULL;

    //二叉查找树插入方法相同
    while (NULL != (*ppRoot))
    {
        ppstPt = *ppNodeTmp;
        if (pNode->key > (*ppNodeTmp)->key)
        {
            ppNodeTmp = &((*ppNodeTmp)->pstRt);
        } 
        else if (pNode->key < (*ppNodeTmp)->key)
        {
            ppNodeTmp = &((*ppNodeTmp)->pstLt);
        }
        else
        {
            return RETURN_ERR;
        }
    }

    *ppNodeTmp = pNode;
    pNode->pstPt = ppstPt;
    pNode->color = RED;
    pNode->pstLt = NULL;
    pNode->pstRt = NULL;

    rb_InsertFixup(pNode, ppRoot);
    
    return RETURN_OK;
}

/*****************************************************************************
函数功能：根据key，查找节点
函数入参：int key  
         rb_t *pRoot 红黑树的根节点
函数出参：无
函数返回值：key对应的节点
其他：
******************************************************************************/
rb_t* rb_Search(int key, rb_t *pRoot)
{
    rb_t *pTmp = pRoot;

    while (NULL != pTmp)
    {
        if (pTmp->key < key)
        {
            pTmp = pTmp->pstLt;
        }
        else if (pTmp->key > key)
        {
            pTmp = pTmp->pstRt;
        }
        else
        {
            return pTmp;
        }
    }

    return NULL;
}

void rb_DeletepDelHaveTwoChildren(rb_t *pDel, rb_t **ppRoot, COLOR *peColor, rb_t **ppDelNxtChild, rb_t **ppDelNxtParent)
{
    rb_t *pTmp = pDel;
    rb_t *pDelNext = NULL;			//待删除节点的后继节点
    rb_t *pDelNxtChild = NULL;	//后继节点的右孩子
    rb_t *pDelNxtParent = NULL;		//后继节点的父节点

    //查找pDel的后继
    pDelNext = pDel->pstRt;
    pDelNext = pDelNext->pstLt;
    while (NULL != pDelNext)
    {
        pDelNext = pDelNext->pstLt;
    }

    pDelNxtChild = pDelNext->pstRt;
    pDelNxtParent = pDelNext->pstPt;
    *peColor = pDelNext->color;

    //后续节点存在右孩子
    if (NULL != pDelNxtChild)
    {
        pDelNxtChild->pstPt = pDelNxtParent;
    }

    if (NULL != pDelNxtParent)
    {
        if (pDelNxtParent->pstLt == pDelNext)
        {
            pDelNxtParent->pstLt = pDelNxtChild;
        }
        else
        {
            pDelNxtParent->pstRt = pDelNxtChild;
        }
    }
    else		//后续节点为空，说明待删除的节点时根节点，需要修改根节点
    {
        *ppRoot = pDelNxtChild;
    }

    if (pDelNext->pstPt == pTmp)
    {
        pDelNxtParent = pDelNext;
    }

    pDelNext->pstPt = pTmp->pstPt;
    pDelNext->color = pTmp->color;
    pDelNext->pstRt = pTmp->pstRt;
    pDelNext->pstLt = pTmp->pstLt;

    if (pTmp->pstPt)
    {
        if (pTmp->pstPt->pstLt == pTmp)
        {
            pTmp->pstPt->pstLt = pDelNext;
        }
        else
        {
            pTmp->pstPt->pstRt = pDelNext;
        }
    }
    else
    {
        *ppRoot = pDel;
    }

    pTmp->pstLt->pstPt = pDel;
    if (pTmp->pstRt)
    {
        pTmp->pstRt->pstPt = pDel;
    }

    *ppDelNxtChild = pDelNxtChild;
    *ppDelNxtParent = pDelNxtParent;

    return;
}

void rb_DeletepDelNoTwoChild(rb_t *pDel, rb_t **ppRoot, COLOR *peColor, rb_t **ppDelNxtChild, rb_t **ppDelNxtParent)
{
    rb_t *pTmp = pDel;
    rb_t *pDelNext = NULL;			//待删除节点的后继节点
    rb_t *pDelNxtChild = NULL;		//后继节点的孩子
    rb_t *pDelNxtParent = NULL;		//后继节点的父节点

    //只有一个孩子的情况
    if (NULL != pDel->pstLt)
    {
        pDelNxtChild = pDel->pstRt;
    }
    else if (NULL != pDel->pstRt)
    {
        pDelNxtChild = pDel->pstLt;
    }

    pDelNxtParent = pDel->pstPt;
    *peColor = pDel->color;

    //修改待删除节点的孩子节点的父节点
    if (pDelNxtChild)
    {
        pDelNxtChild->pstPt = pDelNxtParent;
    }

    if (pDelNxtParent)
    {
        if (pDelNxtParent->pstLt == pDel)
        {
            pDelNxtParent->pstLt = pDelNxtChild;
        }
        else
        {
            pDelNxtParent->pstRt = pDelNxtChild;
        }
    }
    else		//父节点为空说明待删除的节点是根节点，需要修改根节点
    {
        *ppRoot = pDelNxtChild;
    }

    *ppDelNxtChild = pDelNxtChild;
    *ppDelNxtParent = pDelNxtParent;

    return;
}

void rb_DeleteNode(rb_t *pDel, rb_t **ppRoot, COLOR *peColor, rb_t **ppDelNxtChild, rb_t **ppDelNxtParent)
{
    //待删除的节点既有左孩子又有右孩子
    if ((NULL != pDel->pstLt) && (NULL != pDel->pstRt))
    {
        rb_DeletepDelHaveTwoChildren(pDel, ppRoot, peColor, ppDelNxtChild, ppDelNxtParent);
    }
    else
    {
        rb_DeletepDelNoTwoChild(pDel, ppRoot, peColor, ppDelNxtChild, ppDelNxtParent);
    }
    return;
}

void rb_DelFixupLeft(rb_t **ppDelNxtChild, rb_t *pDelNxtParent, rb_t **ppRoot)
{
    rb_t *pTmp;
    rb_t *pTmpLeft;

    pTmp = pDelNxtParent->pstRt;
    if (pTmp->color == RED)   //情况1：待删除节点的兄弟pTmp是红色的  
    {
        pTmp->color = BLK;
        pDelNxtParent->color = RED;   //上俩行，改变颜色，pTmp->黑、待删除的节点的父节点->红。  
        rb_LeftRotate(pDelNxtParent, ppRoot);  //再对待删除的节点的父节点做一次左旋  
        pTmp = pDelNxtParent->pstRt;    //待删除节点的新兄弟new w 是旋转之前w的某个孩子。其实就是左旋后的效果。
    }

    if ((!pTmp->pstLt || pTmp->pstLt->color == BLK) && (!pTmp->pstRt || pTmp->pstRt->color == BLK)) //情况2：x的兄弟w是黑色，且w的俩个孩子也都是黑色的
    {                         
        //由于w和w的俩个孩子都是黑色的，则在x和w上得去掉一黑色，  
        pTmp->color = RED;   //于是，兄弟w变为红色。  
        *ppDelNxtChild = pDelNxtParent;    //p[x]为新结点x  
        pDelNxtParent = (*ppDelNxtChild)->pstPt;  //x<-p[x]  
    }
    else //情况3：x的兄弟w是黑色的， 且，w的左孩子是红色，右孩子为黑色。   
    {  
        if (!pTmp->pstRt || pTmp->pstRt->color == BLK)
        {
            if ((pTmpLeft = pTmp->pstLt))   //w和其左孩子pstLt[w]，颜色交换。  
            {
                pTmpLeft->color = BLK;    //w的左孩子变为由黑->红色  
            }
            pTmp->color = RED;           //w由黑->红  
            rb_RightRotate(pTmp, ppRoot);  //再对w进行右旋，从而红黑性质恢复。  
            pTmp = pDelNxtParent->pstRt;        //变化后的，父结点的右孩子，作为新的兄弟结点  
        }

        //情况4：x的兄弟w是黑色的  
        pTmp->color = pDelNxtParent->color;  //把兄弟节点染成当前节点父节点的颜色。  
        pDelNxtParent->color = BLK;         //把当前节点父节点染成黑色  
        if (pTmp->pstRt)                    //且w的右孩子是红  
        {
            pTmp->pstRt->color = BLK;       //兄弟节点w右孩子染成黑色  
        }
        rb_LeftRotate(pDelNxtParent, ppRoot);   //并再做一次左旋  
        *ppDelNxtChild = *ppRoot;                 //并把x置为根。  
        return;
    }
    return;
}

void rb_DelFixupRight(rb_t **ppDelNxtChild, rb_t *pDelNxtParent, rb_t **ppRoot)
{
    rb_t *pTmp;
    rb_t *pTmpRight;

    pTmp = pDelNxtParent->pstLt;
    if (pTmp->color == RED)
    {
        pTmp->color = BLK;
        pDelNxtParent->color = RED;
        rb_RightRotate(pDelNxtParent, ppRoot);
        pTmp = pDelNxtParent->pstLt;
    }

    if ((!pTmp->pstLt || pTmp->pstLt->color == BLK) && (!pTmp->pstRt || pTmp->pstRt->color == BLK))
    {
        pTmp->color = RED;
        *ppDelNxtChild = pDelNxtParent;
        pDelNxtParent = (*ppDelNxtChild)->pstPt;
    }
    else
    {
        if (!pTmp->pstLt || pTmp->pstLt->color == BLK)
        {
            if ((pTmpRight = pTmp->pstRt))
            {
                pTmpRight->color = BLK;
            }
            pTmp->color = RED;
            rb_LeftRotate(pTmp, ppRoot);
            pTmp = pDelNxtParent->pstLt;
        }
        pTmp->color = pDelNxtParent->color;
        pDelNxtParent->color = BLK;
        if (pTmp->pstLt)
        {
            pTmp->pstLt->color = BLK;
        }
        rb_RightRotate(pDelNxtParent, ppRoot);
        *ppDelNxtChild = *ppRoot;
        return;
    }
    return;
}

void rb_DelFixup(rb_t *pDelNxtChild, rb_t *pDelNxtParent, rb_t **ppRoot)
{
    while ((!pDelNxtChild || pDelNxtChild->color == BLK) && pDelNxtChild != *ppRoot)
    {
        if (pDelNxtParent->pstLt == pDelNxtChild)
        {
            rb_DelFixupLeft(&pDelNxtChild, pDelNxtParent, ppRoot);
        }
        else
        {
            rb_DelFixupRight(&pDelNxtChild, pDelNxtParent, ppRoot);
        }
    }

    if (pDelNxtChild)
    {
        pDelNxtChild->color = BLK;
    }
    return;
}

/*****************************************************************************
函数功能：将一个节点从红黑树中删除（只是将节点从红黑树中摘掉，节点的内存不会再本函数中释放）
函数入参：rb_t *pDel
         rb_t **ppRoot
函数出参：无
函数返回值：无
特别说明：pDel的内存需要在调用此函数之后，手动释放
******************************************************************************/
void rb_Delete(rb_t *pDel, rb_t **ppRoot)
{
    COLOR color;
    rb_t *pDelNxtChild = NULL;
    rb_t *pDelNxtParent = NULL;

    //将待删除的节点从红黑树中摘掉
    rb_DeleteNode(pDel, ppRoot, &color, &pDelNxtChild, &pDelNxtParent);

    if (color == BLK)
    {
        rb_DelFixup(pDelNxtChild, pDelNxtParent, ppRoot); //调用rb_erase_rebalance来恢复红黑树性质
    }

    return;
}