紅黑樹的基本操作

rbTree.h

#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

rbTree.c

#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;
}