红黑二叉查找树

#include <stdio.h>

//红黑二叉查找树

//红黑二叉查找树是一种平衡二叉树,是基于2-3查找树的基础上演变的
//这里不对2-3查找树的算法进行描述,感兴趣的朋友可以自行了解一下
//2-3查找树的实现原理.
//
//在2-3查找树算法中,难点就是3-Node类型(含有2个key,3个子节点的节点类型)
//的节点数据结构处理,那么在红黑二叉查找树中这种3-Node类型节点演变为如
//下表示方式.非常巧妙的在原普通二叉数结点数据结构中增加了一个颜色的成
//员变量,就使得红黑二叉查找树可以替换为2-3查找树.当一个节点的左子节点
//的颜色为红色时，该节点即为2-3查找树中的3-Node类型节点，此时可以将红
//色的子节点理解为和父节点是平行的，一体的，比如当前示例中将Node1与
//Node2合并成一个节点，该节点相当于2-3查找树中的3-Node类型节点，即含有
//2个key，3个子节点。
//
//  ** 相当于2-3查找树中的3-Node类型节点:
//                Node1(Black)
//              --------------
//             /              \
//        Node2(Red)      NULL(Black)
//       -----------
//      /           \
//  NULL(Black)  NULL(Black)
//  
//  ** 相当于2-3查找树中的2-Node类型节点:
//                Node1(Black)
//              --------------
//             /              \
//        Node2(Black)      NULL(Black)
//       -----------
//      /           \
//  NULL(Black)  NULL(Black)
//
//红黑二叉查找树的定义如下：
//1 根节点必须为黑色
//2 红链接均为左链接
//3 没有任何一个结点同时和两条红链接相连
//4 该树是完美黑色平衡的，即任意空链接到根结点的路径上的黑链接数量
//  相同。
//备注： 空节点的颜色为黑色。

typedef struct _NODE
{
    int             key;
    int             val;
    int             n;
    int             isRed;
    struct _NODE*   left;
    struct _NODE*   right;
} NODE;

enum
{
    FALSE,
    TRUE
} BOOL_TYPE;

NODE* root = NULL;

int isRed(NODE* pNode)
{
    if ( pNode == NULL )
    {
        return FALSE; 
    }

    return pNode->isRed;
}

int size(NODE* pNode)
{
    if ( pNode == NULL )
    {
        return 0; 
    }

    return pNode->n;
}

int compare(int key1, int key2)
{
    if ( key1 < key2 )
    {
        return -1; 
    }

    if ( key1 > key2 )
    {
        return 1; 
    }

    return 0;
}

NODE* newNode(int key, int val, int isRed)
{
    NODE* pNew = (NODE*)malloc(sizeof(NODE));
    if ( pNew == NULL )
    {
        return NULL; 
    }

    memset(pNew, 0, sizeof(NODE));
    pNew->key = key;
    pNew->val = val;
    pNew->isRed = isRed;
    pNew->n = 1;

    return pNew;
}

//树的左旋：
//1 首先把该节点的右子节点进行提升
//2 之后把该节点进行下降(由于右子节点提升，该节点下降，所以此刻该节点的等级
//  仅仅和新位置的右子点的子节点平行，即当前右子节点已经可以做为被旋转节点
//  的父节点了)
//3 将该节点的右链接指向原右子节点的左树
//4 将原右子节点的左链接指向该节点
// （3、4步之所以这样操作，主要是为了满足二叉树的特性，每个节点最多有两个子
// 节点，同时左子节点一定小于该节点，右子节点一定大于该节点。）
//5 可以看到当前仅被旋转节点和该节点的右子节点做了变化，所以仅对这两个节点
//  的相关信息进行更新。
//
//示例：对5节点进行左旋
//
//     旋转前
//       5
//     /   \
//    3     7
//   / \   / \
//  2   4 6   8
//
//     旋转后
//        7
//       / \
//      5   8
//     / \
//    3   6
//   / \
//  2   4 
//
NODE* rotateLeft(NODE* pNode)
{
    NODE* pReturnNode = NULL;

    if ( pNode == NULL || pNode->right == NULL )
    {
        return pNode; 
    }

    pReturnNode = pNode->right;

    pNode->right = pNode->right->left; 
    pReturnNode->left = pNode;

    //这里将新提升的结点的颜色设置为原来结点的颜色，目地就是为了保证
    //原结点上面的父结点相关路径的颜色不变。
    //当然，此时被旋转的节点进行了下降，此时该节点的颜色修改为红色，
    //如果之前节点是黑色，则此时变成红色，符合红黑二叉查找树的特性。
    //但如果之前是红色，此时变成红色，则违背了红黑二叉查找树的特性
    //(没有任何一个结点同时和两条红链接相连)，这时候就需要不断进行向
    //上进行父节点的颜色调整，使得最终端满足红黑树的特性。从这里也可
    //以看出，其实原来为黑，后面变红，就相当于2-3查找树算法中的
    //2-Node类型节点变成3-Node类型节点;而原来为红，后面变红，就相当
    //于2-3查找树算法中将3-Node类型节点变成临时4-Node类型节点，之后
    //将4-Node类型节点拆分为3个2-Node类型节点，并将中间的2-Node节点
    //继续向上与父节点进行合并
    pReturnNode->isRed = pNode->isRed;
    pNode->isRed = TRUE;

    pReturnNode->n = pNode->n;
    pNode->n = size(pNode->left) + size(pNode->right) + 1;

    return pReturnNode;
}

//树的右旋：
//同左旋原理相同，这里仅给出示例
//
//示例：对5节点进行右旋
//
//     旋转前
//       5
//     /   \
//    3     7
//   / \   / \
//  2   4 6   8
//
//     旋转后
//        3
//       / \
//      2   5
//         / \
//        4   7
//           / \
//          6   8 
//
NODE* rotateRight(NODE* pNode)
{
    NODE* pReturnNode = NULL;

    if ( pNode == NULL || pNode->left == NULL )
    {
        return pReturnNode; 
    }

    pReturnNode = pNode->left;

    pNode->left = pNode->left->right;
    pReturnNode->right = pNode;

    pReturnNode->isRed = pNode->isRed;
    pNode->isRed = TRUE;

    pReturnNode->n = pNode->n;
    pNode->n = size(pNode->left) + size(pNode->right) + 1;

    return pReturnNode;
}

void flipColors(NODE* pNode)
{
    if ( pNode == NULL || pNode->left == NULL || pNode->right == NULL )
    {
        return; 
    }

    pNode->isRed = TRUE;
    pNode->left->isRed = FALSE;
    pNode->right->isRed = FALSE;
}

//红黑二叉查找树的插入
//这里同2-3查找树算法类似，主要考虑5种情况：
//1 插入到2-Node类型节点的左侧
//2 插入到2-Node类型节点的右侧
//3 插入到3-Node类型节点的左侧
//4 插入到3-Node类型节点的中侧
//5 插入到3-Node类型节点的右侧
//
// * 插入到2-Node类型节点的左侧
//   如下示例所示，将2插入到4节点的左侧，新建立的节点都是红色,此时
//   满足红黑树的特性，所以后继不需要在做其它处理。
//
// 插入前：
//    4(B)
//    /  \
//  NULL  6(B)
//
// 插入后：
//    4(B)
//    /  \
//  2(R) 6(B)
//
// * 插入到2-Node类型节点的右侧
// 如下示例所示，将6插入到4节点的右侧，新建立的节点是红色，此时
// 违背了红黑树的特性，不能出现红色的右链接，为了恢复红黑树的特性，
// 这里需要再将4节点进行左旋转，可以参见树左旋的介绍，这样红黑
// 树的特性又满足了。
//
// 插入前：
//    4(B)
//    /  \
//  5(B)  NULL
//
// 插入后：
//    4(B)
//    /  \
//  5(B) 6(R)
//
// * 插入到3-Node类型节点的右侧
// 先介绍这种情况，因为这种情况在3-Node类型节点的几种插入算法中最简单。
// 如下示例所示，将8插入到7节点的右侧，新建立的节点为红色，由于这里是
// 将新节点插入到3-Node类型节点中，所以不能依照2-Node类型算法中的插入
// 右侧进行分析，3-Node类型节点的插入有自己单独的算法，在插入3-Node类
// 型节点右侧的情况下，将7的左、右两个子节点5和8的颜色全部变成黑色，之
// 后7节点自身的颜色变成红色，这时候关键就看7这个节点，如果7节点也是一棵
// 子树，则由于7节点颜色发生了变化，所以需要检测7节点上面的父节点情况，
// 一般情况下父节点会有3种情况
// a 如果父节点不存在，则7就是最根节点，则根节点7为了满足红黑树的特性，
//   必须变成黑色
// b 如果父节点为2-Node类型节点，则按照上面2-Node的两种方法进行调用父
//   节点也当前子树根节点的关系即可
// c 如果父节点为3-Node类型节点，则同当前3-Node类型节点的方法相同，调
//   整父节点与当前子树根节点的关系，采用这种循环处理父节点的方式，
//   直到上升调整到父节点为空或为2-Node类型则停止。
// (此时的算法已经十分类似向2-3查找树算法中，向3-Node类型节点插入相似)
//
// 插入前：
//      7(B)
//      /  \
//    5(R)  NULL
//    /  \
//  3(B) 6(B)
//
// 插入后：
//       7(B)
//      /   \
//    5(R)  8(R)
//    /  \
//  3(B) 6(B)
//
// * 插入到3-Node类型节点的左侧
// 如示例所示，将节点3插入到3-Node类型节点左侧，新建的节点是红色，
// 插入后，出现7的左节点为红色，同时左节点的左节点也是红色，此时不
// 能满足红黑树的特性，之后将节点7进行右旋，仔细查看右旋后的情况，
// 和之前"插入到3-Node类型节点的右侧"后的情况竟然相同了，此时参数
// 上述方法，将两个子节点变成黑色，当前子树的根节点变成红色即可。
//
// 插入前：
//       7(B)
//      /   \
//    5(R)  8(B)
//    /  \
//       6(B)
//
// 插入后：
//       7(B)
//      /   \
//    5(R)  8(B)
//    /  \
//  3(R) 6(B)
//
// 右旋后：
//     5(B)
//    /    \
//  3(R)   7(R)
//         /  \
//       6(B) 8(B)
//
// * 插入到3-Node类型节点的中侧
// 最后一种插入情况了，放最后是因为这种情况最麻烦
// 如示例所示,将6插入到3-Node类型节点的中侧，插入后，新节点为红色，
// 不满足红黑树中右链接不能出现红色的特性，此时将5节点进行左旋转，旋转
// 后仔细观察，其中这时候出现左节点、以及左节点的左节点为红色，这个现象
// 同之前的"插入到3-Node类型节点的左侧"的相同了，所以采用那种情况的解决
// 方法，将当前节点7再进行一次右旋转，这时候又出现了左、右子节点都为红
// 色的现象，采用之前的方法，将两个子节点变成黑色，子树的根节点变成红
// 色。
//
// 插入前：
//       7(B)
//      /   \
//    5(R)  8(B)
//    /  \
//  3(B)
//
// 插入后：
//       7(B)
//      /   \
//    5(R)  8(B)
//    /  \
//  3(B) 6(R)
//
// 对5节点进行左旋后：
//        7(B)
//       /   \
//     6(R)  8(B)
//     /
//   5(R)
//    /
//  3(B)
//
// 对7节点进行右旋后：
//        6(B)
//       /   \
//     5(R)  7(R)
//     /       \
//   3(B)     8(B)
//
NODE* _put(NODE* pNode, int key, int val)
{
    int cmp = 0;

    if ( pNode == NULL )
    {
        return newNode(key, val, TRUE); 
    }

    cmp = compare(key, pNode->key);

    if ( cmp < 0 )
    {
        pNode->left = _put(pNode->left, key, val);
    }
    else if ( cmp > 0 )
    {
        pNode->right = _put(pNode->right, key, val);
    }
    else
    {
        pNode->val = val;    
    }

    //这里处理“插入到3-Node类型节点的中侧"时候的情况，注意空节点也是黑色。
    //同时这里也处理"插入到2-Node类型节点的右侧"的情况。
    if ( isRed(pNode->right) && !isRed(pNode->left) )
    {
        pNode = rotateLeft(pNode); 
    }

    //这里处理"插入到3-Node类型节点的左侧"时候的情况。
    if ( isRed(pNode->left) && isRed(pNode->left->left) )
    {
        pNode = rotateRight(pNode); 
    }

    //这里处理“插入到3-Node类型节点的右侧"时候的情况。
    if ( isRed(pNode->left) && isRed(pNode->right) )
    {
        flipColors(pNode);
    }

    //提示：
    //在上述3种颜色处理流程的顺序不能出现颠倒，可以参考函数头部中注释的讲解，
    //之所以这里存在顺序的要求，是因为"3-Node类型中侧"处理情况可以演变成"3-Node
    //类型左侧"处理情况，而"3-Node类型左侧"处理情况又可以演变成"3-Node类型右侧"
    //处理情况，所以这里需要存在这种顺序关系。

    pNode->n = size(pNode->left) + size(pNode->right) + 1;

    return pNode;
}

void put(int key, int val)
{
    root = _put(root, key, val);
    root->isRed = FALSE;
}

//中序打印
//这里同普通二叉树的算法相同，不再进行描述
void showTree(NODE* pNode, int level)
{
    if ( pNode == NULL )
    {
        return; 
    }

    showTree(pNode->left, level + 1);
    printf("[level %3d] key %d, val %d, %s\r\n", \
            level,  pNode->key, pNode->val, pNode->isRed ? "RED" : "BLACK");
    showTree(pNode->right, level + 1);
}

//查找
//这里同普通二叉树查找算法相同，不再进行描述
int _get(NODE *pNode, int key)
{
    if ( pNode == NULL )
    {
        return -1; 
    }

    if ( key < pNode->key )
    {
        return _get(pNode->left, key); 
    }
    else if ( key > pNode->key )
    {
        return _get(pNode->right, key); 
    }

    return pNode->val;
}

int get(int key)
{
    return _get(root, key);
}

int main()
{
    int i = 0;
    int keys[] = {2,  4,  6,  1,  9,  4,  9,  8,  7,  2};
    int vals[] = {21, 41, 61, 11, 91, 42, 92, 81, 71, 22};

    for ( i = 0; i < (sizeof(keys) / sizeof(keys[0])); i++ )
    {
        put(keys[i], vals[i]);    
    }

    showTree(root, 0);

    printf("find key %d, get val %d\r\n", 4,  get(4));
    printf("find key %d, get val %d\r\n", 9,  get(9));
    printf("find key %d, get val %d\r\n", 3,  get(3));

    return 0;
}