C++實現一個簡單的紅黑樹（RB_TREE）

紅黑樹遵守的規則：一頭（根）一腳（葉子節點）黑（黑色），黑同（從任意節點開始至NULL結點的黑色節點的個數相等）紅不連（紅色節點不能相連）

以下是紅黑樹的插入和旋轉函數的簡易實現：

#include <iostream>
#include <assert.h>
#include <string.h>

using namespace std;

typedef int Type;
typedef enum{RED=0, BLACK}Color;    //節點的顏色

typedef struct Node{
    Color color;   //節點顏色
    Type key;      //節點的值
    struct Node *left, *right, *parent;  //左，右及父節點指針
}*PNode;   //節點指針


//樹的結構定義
typedef struct RB_TREE{
    Node *root;   //根結點
    Node *Nil;    //一個實現技巧，用於判斷節點是否爲空等
}RB_TREE;


Node *Buynode()
{
    Node *p = new Node;
    assert(p != NULL);
    memset(p, 0, sizeof(Node));

    return p;
}


//構造一個空樹
void InitTree(RB_TREE &t)
{
    t.Nil = Buynode();
    t.root = t.Nil;
    t.Nil->color = BLACK;
    t.Nil->key = -1;
}


//左旋
void RotateLeft(RB_TREE &t, Node *p)
{
    Node *s = p->right;  //s爲不平衡節點p的右樹
    p->right = s->left;   //將s的左樹掛接至p的右樹
    if(s->left != t.Nil){  //若s的左樹不爲空,就改變其左樹的父節點
        s->left->parent = p;
    }
    s->parent = p->parent;  //改變s的父節點

    if(p->parent == t.Nil){  //說明爲p爲根結點，旋轉之後s爲根結點
        t.root = s;
    }else if(p = p->parent->left){  //p之前位於左樹就將s作爲新的左樹
        p->parent->left = s;
    }else{
        p->parent->right = s;   //否則s爲p的父節點的右樹
    }

    s->left = p;   //p作s的左樹
    p->parent = s;  //更改p的父節點
}


//右旋
void RotateRight(RB_TREE &t, Node *p)
{
    Node *s = p->left;
    p->left = s->right;
    if(s->right != t.Nil){
        s->right->parent = p;
    }
    s->parent = p->parent;

    if(p->parent == t.Nil){
        t.root = s;
    }else if(p = p->parent->left){
        p->parent->left = s;
    }else{
        p->parent->right = s;
    }

    s->right = p;
    p->parent = s;
}

//調整樹的平衡
void Insert_Fixup(RB_TREE &t, Node *z)
{
    Node *y;

    while(z->parent->color == RED){   //紅紅相連不平衡
        if(z->parent == z->parent->parent->left){   //左側插入
            y = z->parent->parent->right;   //y爲插入節點的伯父節點
            if(y->color == RED){
                z->parent->color = BLACK;
                z->parent->parent->color = RED;

                y->color = BLACK;
                z = z->parent->parent;
                continue;
            }else if(z == z->parent->right){  //左側的內側插入
                z = z->parent;
                RotateLeft(t, z);    //左旋
            }

            z->parent->color = BLACK;
            z->parent->parent->color = RED;
            RotateRight(t, z->parent->parent);   //右旋
        }else{   //右側插入
            y = z->parent->parent->left;
            if(y->color == RED){
                z->parent->color = BLACK;
                z->parent->parent->color = RED;

                y->color = BLACK;
                z = z->parent->parent;
                continue;
            }else if(z == z->parent->left){   //右側的內側插入
                z = z->parent;
                RotateRight(t, z);
            }

            z->parent->color = BLACK;
            z->parent->parent->color = RED;
            RotateLeft(t, z->parent->parent);
        }
    }

    t.root->color = BLACK;
}


bool Insert(RB_TREE &t, Type x)
{
    Node *p = t.Nil;
    Node *s = t.root;

    //找到合適的插入位置
    while(s != t.Nil){
        p = s;
        if(x == s->key){
            return false;
        }else if(x < s->key){
            s = s->left;
        }else{
            s = s->right;
        }
    }

    //構造節點
    Node *q = Buynode();
    q->key = x;
    q->parent = p;

    //將節點插入樹中的合適位置
    if(p == t.Nil){   //則說明該樹之前沒有節點，則此節點爲其根結點
        t.root = q;
    }else if(x < p->key){  //插入節點在節點p的左子樹
        p->left = q;
    }else{    //插入節點在右子樹
        p->right = q;
    }

    //設置插入節點的信息
    q->left = q->right = t.Nil;
    q->color = RED;
    //調整樹的平衡
    Insert_Fixup(t, q);
    return true;
}

int main()
{
    int ar[] = {100, 40, 6};
    RB_TREE rb;
    InitTree(rb);

    int n = sizeof(ar) / sizeof(int);
    for(int i = 0; i < n; ++i){
        Insert(rb, ar[i]);
    }

    return 0;
}