学步园

原文来自：http://liyiwen.iteye.com/blog/345799

// -------------------------------------------------------  

// FileName : BinarySearchTree.cpp  
// Date     : 2009-1-14   
// Summary  : 红黑树的简单实现，节点中除了红黑树结构本身需要的信息之外，只  
//            包含一个int域作为节点的key  
// CopyLeft : 李诒雯  
// -------------------------------------------------------  
  

#include <stdio.h>  

#include <stdlib.h>

  
// 定义节点颜色  
enum COLOR {  
    BLACK = 0,  
    RED  
};  
  
// 红黑树节点  
typedef struct tagRBTreeNode {  
    int key;  
    struct tagRBTreeNode *left;  
    struct tagRBTreeNode *right;  
    struct tagRBTreeNode *parent;  
    unsigned char color;  
}RBTreeNode, *PRBTreeNode;  
  
// 红黑树，包含一个指向根节点的指针  
typedef struct tagRBTree {  
    PRBTreeNode root;     
}RBTree, *PRBTree;  
  
// 红黑树的NIL节点  
static RBTreeNode NIL = {0, 0, 0, 0, BLACK};  
  
#define PNIL (&NIL)   // NIL节点地址  
  
// 生成一个红黑树节点，缺省的节点颜色是红黑。  
PRBTreeNode CreateRBTreeNode(int keyVal)  
{  
    PRBTreeNode p = (PRBTreeNode)malloc(sizeof(RBTreeNode));  
    p->color = RED;  
    p->left = PNIL;  
    p->right = PNIL;  
    p->parent = PNIL;  
    p->key = keyVal;  
    return p;  
}  
  
// 释放一个节点的内存  
void DeleteRBTreeNode(PRBTreeNode pDel)  
{  
    if (PNIL != pDel) {  
        free(pDel);  
    }     
}  
  
void InitRBTree(PRBTree pTree) // 初始化一棵红黑树  
{  
    pTree->root = PNIL;  
}  
  
// 顺序遍历一棵红黑树，并对每个节点执行指定的操作  
void INORDER_RBTREE_WALK(PRBTreeNode pRoot, void (*Action)(PRBTreeNode ))  
{  
    // 就是一个二叉树的中序遍历  
    if (PNIL != pRoot) {  
        INORDER_RBTREE_WALK(pRoot->left, Action);  
        Action(pRoot);  
        INORDER_RBTREE_WALK(pRoot->right, Action);  
    }  
}  
  
// 查找指定键值的节点  
PRBTreeNode RBTREE_SEARCH(PRBTreeNode pRoot, int val_key)  
{  
    while ((PNIL != pRoot) && (val_key != pRoot->key)) {  
        pRoot = (val_key < pRoot->key ? pRoot->left : pRoot->right);      
    }  
    return pRoot;  
}  
  
// 查找最大键值节点  
PRBTreeNode RBTREE_MAX(PRBTreeNode pRoot)  
{  
    while (PNIL != pRoot->right) {  
        pRoot = pRoot->right;      
    }  
    return pRoot;  
}  
  
// 查找最小键值节点  
PRBTreeNode RBTREE_MIN(PRBTreeNode pRoot)  
{  
    while (PNIL != pRoot->left) {  
        pRoot = pRoot->left;   
    }  
    return pRoot;  
}  
  
// 查找指定节点的后继节点  
PRBTreeNode RBTREE_SUCCESSOR(PRBTreeNode pRoot)  
{  
    if (PNIL != pRoot->right) {  
        return RBTREE_MIN(pRoot->right);  
    }  
    PRBTreeNode pParent = pRoot->parent;  
    while((PNIL != pParent) && (pRoot == pParent->right)) {  
        pRoot = pParent;  
        pParent = pRoot->parent;           
    }   
    return pParent;  
}  
  
// 查找指定节点的前趋节点  
PRBTreeNode RBTREE_PREDECESSOR(PRBTreeNode pRoot)  
{  
    if (PNIL != pRoot->left) {  
        return RBTREE_MAX(pRoot->left);  
    }  
    PRBTreeNode pParent = pRoot->parent;  
    while((PNIL != pParent) && (pRoot == pRoot->left)) {  
        pRoot = pParent;  
        pParent = pRoot->parent;           
    }   
    return pParent;  
}  
  
// 对指定节点进行左旋  
static void LEFT_ROTATION(PRBTree pTree, PRBTreeNode pNode)  
{  
    PRBTreeNode pRight = pNode->right;  
    pRight->parent = pNode->parent;  
    if(pNode->parent->left == pNode){  
        pRight->parent->left = pRight;  
    }  
    else {  
        pRight->parent->right = pRight;  
    }  
  
    if (PNIL == pRight->parent) {  
        pTree->root = pRight;  
    }  
    if (PNIL != pRight->left) {  
        pRight->left->parent = pNode;  
    }  
    pNode->right = pRight->left;  
    pNode->parent = pRight;  
    pRight->left = pNode;  
}  
  
// 对指定节点进行右旋  
static void RIGHT_ROTATION(PRBTree pTree, PRBTreeNode pNode)  
{  
    PRBTreeNode pLeft = pNode->left;  
    pLeft->parent = pNode->parent;  
    if(pNode->parent->left == pNode){  
        pLeft->parent->left = pLeft;  
    }  
    else {  
        pLeft->parent->right = pLeft;  
    }  
  
    if (PNIL == pLeft->parent) {  
        pTree->root = pLeft;  
    }  
    if (PNIL != pLeft->right) {  
        pLeft->right->parent = pNode;  
    }     
    pNode->left = pLeft->right;  
    pNode->parent = pLeft;  
    pLeft->right = pNode;  
}  
  
// 在红黑树插入动作后进行红黑性质恢复  
static void FIX_UP_RB_AFTER_INSERT(PRBTree pTree, PRBTreeNode pNode)  
{  
    while(RED == pNode->parent->color) {  
        if (pNode->parent == pNode->parent->parent->left){  
            PRBTreeNode pUncleNode = pNode->parent->parent->right;  
            if (RED == pUncleNode->color){  
                pNode->parent->color = BLACK;  
                pUncleNode->color = BLACK;  
                pNode = pUncleNode->parent;  
                pNode->color = RED;  
            }  
            else   
            {  
                if (pNode == pNode->parent->right) {  
                    pNode = pNode->parent;  
                    LEFT_ROTATION(pTree, pNode);  
                }  
                pNode->parent->color = BLACK;  
                PRBTreeNode TempP = pNode->parent->parent;  
                TempP->color = RED;  
                RIGHT_ROTATION(pTree, TempP);  
            }  
        }  
        else{  
            PRBTreeNode pUncleNode = pNode->parent->parent->left;  
            if (RED == pUncleNode->color){  
                pNode->parent->color = BLACK;  
                pUncleNode->color = BLACK;  
                pNode = pUncleNode->parent;  
                pNode->color = RED;  
            }  
            else   
            {  
                if (pNode == pNode->parent->left) {  
                    pNode = pNode->parent;  
                    RIGHT_ROTATION(pTree, pNode);  
                }  
                pNode->parent->color = BLACK;  
                PRBTreeNode TempP = pNode->parent->parent;  
                TempP->color = RED;  
                LEFT_ROTATION(pTree, TempP);  
            }  
        }  
    }  
    pTree->root->color = BLACK;  
}  
  
// 红黑树的节点插入  
void RBTREE_INSERT(PRBTree pTree, PRBTreeNode pNode)  
{  
    PRBTreeNode pNodePosition = pTree->root;  
    PRBTreeNode pParent = PNIL;  
  
    while (PNIL != pNodePosition) {  
        pParent = pNodePosition;  
        pNodePosition = (pNode->key > pNodePosition->key ?   
            pNodePosition->right : pNodePosition->left);  
    }  
      
    pNode->parent = pParent;  
    if (PNIL == pParent) {  
        pTree->root = pNode;       
    }  
    else {  
        if (pNode->key > pParent->key) {  
            pParent->right = pNode;  
        }  
        else {  
            pParent->left = pNode;  
        }  
    }  
    FIX_UP_RB_AFTER_INSERT(pTree, pNode);  
}  
  
// 在红黑树删除动作后进行红黑性质恢复  
static void FIX_UP_RB_AFTER_DELETE(PRBTree pTree, PRBTreeNode pNode)  
{  
    while ((pNode->color == BLACK) && (pNode != pTree->root)) {  
        if (pNode == pNode->parent->left){  
            PRBTreeNode pBrother = pNode->parent->right;  
            if (RED == pBrother->color) {  
                pBrother->color = BLACK;  
                pNode->parent->color = RED;  
                LEFT_ROTATION(pTree, pNode->parent);  
                pBrother = pNode->parent->right;  
            }  
            if ((BLACK == pBrother->right->color) && (BLACK == pBrother->left->color)) {  
                pBrother->color = RED;  
                pNode = pNode->parent;  
            }  
            else {  
                if (BLACK == pBrother->right->color) {  
                    pBrother->color = RED;  
                    pBrother->left->color = BLACK;  
                    RIGHT_ROTATION(pTree, pBrother);  
                    pBrother = pNode->parent->right;  
                }  
                pBrother->right->color = BLACK;  
                pBrother->color = pNode->parent->color;                  
                pNode->parent->color = BLACK;  
                LEFT_ROTATION(pTree, pNode->parent);  
                pNode = pTree->root;  
            }  
        }  
        else{  
            PRBTreeNode pBrother = pNode->parent->left;  
            if (RED == pBrother->color) {  
                pBrother->color = BLACK;  
                pNode->parent->color = RED;  
                RIGHT_ROTATION(pTree, pNode->parent);  
                pBrother = pNode->parent->left;  
            }  
            if ((BLACK == pBrother->left->color) && (BLACK == pBrother->right->color)) {  
                pBrother->color = RED;  
                pNode = pNode->parent;  
            }  
            else {  
                if (BLACK == pBrother->left->color) {  
                    pBrother->color = RED;  
                    pBrother->right->color = BLACK;  
                    LEFT_ROTATION(pTree, pBrother);  
                    pBrother = pNode->parent->left;  
                }  
                pBrother->left->color = BLACK;  
                pBrother->color = pNode->parent->color;                  
                pNode->parent->color = BLACK;  
                RIGHT_ROTATION(pTree, pNode->parent);  
                pNode = pTree->root;  
            }  
        }  
    }  
    pNode->color = BLACK;  
}  
  
// 红黑树的节点删除  
PRBTreeNode RBTREE_DELETE(PRBTree pTree, PRBTreeNode pDel)  
{  
    PRBTreeNode pRelDel;  
    if ((PNIL == pDel->left) || (PNIL == pDel->right)) {  
        pRelDel = pDel;  
    }  
    else { //  
        pRelDel = RBTREE_SUCCESSOR(pDel);  
    }  
  
    PRBTreeNode pChildOfDel;  
    pChildOfDel = (pRelDel->left!= PNIL ? pRelDel->left : pRelDel->right);  
  
    //  
    pChildOfDel->parent = pRelDel->parent;  
      
    if (PNIL == pRelDel->parent) {  
        pTree->root = pChildOfDel;  
    }  
    else {  
        if (pRelDel == pRelDel->parent->left){  
            pRelDel->parent->left = pChildOfDel;  
        }  
        else {  
            pRelDel->parent->right = pChildOfDel;  
        }  
    }  
  
    if (pRelDel != pDel) {  
        pDel->key = pRelDel->key;  
    }  
  
    if (pRelDel->color == BLACK) {  
        FIX_UP_RB_AFTER_DELETE(pTree, pChildOfDel);  
    }  
    return pRelDel;  
}  
  
/*=============================================*/  
// 测试用函数，打印节点键值。可以传给遍历函数作为第二个参数  
void PrintNode(PRBTreeNode p)  
{  
    if (NULL != p) {  
        printf("%2d ", p->key);  
    }  
    else {  
        printf("NULL!");  
    }  
}  
  
// 打印出红黑树每个节点的各项值。  
void RBTREE_STRUCTURE(PRBTreeNode pRoot)  
{  
#define GetKey(p) (p==PNIL?-1:p->key)  
    //  
    if (PNIL != pRoot) {  
        printf("KeyValue = %d, color = %s, Left = %d, Right = %d, Parent = %d\n",  
            pRoot->key, pRoot->color==BLACK?"BLACK":"RED", GetKey(pRoot->left),   
            GetKey(pRoot->right), GetKey(pRoot->parent));  
        RBTREE_STRUCTURE(pRoot->left);  
        RBTREE_STRUCTURE(pRoot->right);  
    }  
}  
  
/*=============================================*/  
int _tmain(int argc, _TCHAR* argv[])  
{  
    RBTree bsTree;  
    InitRBTree(&bsTree);  
  
#define TEST_NUM_COUNT (10)  
    int values[TEST_NUM_COUNT];  
  
    for (int i=0; i<TEST_NUM_COUNT; ++i) {  
        values[i] = rand()%(TEST_NUM_COUNT*10);  
        PRBTreeNode p = CreateRBTreeNode(values[i]);  
        RBTREE_INSERT(&bsTree, p);  
      
          
        printf("After Insert %2d:  ", values[i]);  
        INORDER_RBTREE_WALK(bsTree.root, PrintNode);      
        printf("RB-Tree Nodes Information\n");  
          
        RBTREE_STRUCTURE(bsTree.root);  
        printf("-------------------------------\n");  
    }  
  
    printf("###############################\n");  
    printf("===============================\n");  
  
    for (int i=0; i<TEST_NUM_COUNT; ++i) {  
        PRBTreeNode pDel = RBTREE_DELETE(&bsTree, RBTREE_SEARCH(bsTree.root, values[i]));  
        DeleteRBTreeNode(pDel);  
        printf("After Delete %2d:  ", values[i]);  
        INORDER_RBTREE_WALK(bsTree.root, PrintNode);  
        printf("RB-Tree Nodes Information\n");  
        RBTREE_STRUCTURE(bsTree.root);  
        printf("-------------------------------\n");  
          
    }  
  
    return 0;  
}