红黑树实现及性质验证

[eiclpy]

#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <cassert>
#include <random>
#include <ctime>

using namespace std;

default_random_engine RandInt(time(0));

#define BLACK 0
#define RED 1

template<typename Key, typename Mapped>
class RedBlackTree {
#ifdef DEBUG
public:
#endif
    struct RedBlackTreeNode {
        RedBlackTreeNode *pChildNode[2], *pFatherNode;
        bool bNodeColor;
        Key key;
        Mapped mapped;
    };

    typedef RedBlackTreeNode *RBTree;

    RBTree NullNode, TreeRoot;
    RBTree *DeleteStack;
    int DeleteStackTop;

    void ClearDeleteStack() {
        while (DeleteStackTop)free(DeleteStack[DeleteStackTop--]);
    }

    RBTree CreateNode(Key k, Mapped val, RBTree fa) {
        RBTree ret = (RBTree) malloc(sizeof(RedBlackTreeNode));
        assert(ret);
        ret->pChildNode[0] = ret->pChildNode[1] = NullNode;
        ret->bNodeColor = RED;
        ret->pFatherNode = fa;
        ret->key = k;
        ret->mapped = val;
        return ret;
    }

    void _clear(RBTree pCurrentNode){
        if(pCurrentNode->pChildNode[0]!=NullNode)_clear(pCurrentNode->pChildNode[0]);
        if(pCurrentNode->pChildNode[1]!=NullNode)_clear(pCurrentNode->pChildNode[1]);
        free(pCurrentNode);
    }

    void Rotate(RBTree pCurrentNode, bool Direction) {
        RBTree t = pCurrentNode->pChildNode[!Direction];
        pCurrentNode->pChildNode[!Direction] = t->pChildNode[Direction];
        t->pChildNode[Direction]->pFatherNode = pCurrentNode;
        t->pChildNode[Direction] = pCurrentNode;
        t->pFatherNode = pCurrentNode->pFatherNode;
        if (pCurrentNode->pFatherNode != nullptr)
            pCurrentNode->pFatherNode->pChildNode[pCurrentNode->pFatherNode->pChildNode[0] != pCurrentNode] = t;
        else TreeRoot = t;
        pCurrentNode->pFatherNode = t;
    }

    RBTree _insert(RBTree &pCurrentNode, Key key, Mapped mapped, RBTree fa) {
        if (pCurrentNode == NullNode)return pCurrentNode = CreateNode(key, mapped, fa);
        if (pCurrentNode->key == key)return nullptr;
        return _insert(pCurrentNode->pChildNode[pCurrentNode->key < key], key, mapped, pCurrentNode);
    }

    RBTree _delete(RBTree &pCurrentNode, int key) {
        if (pCurrentNode == NullNode)return nullptr;
        if (pCurrentNode->key == key) {
            RBTree pChildNode = nullptr;
            if (pCurrentNode->pChildNode[0] == NullNode)pChildNode = pCurrentNode->pChildNode[1];
            if (pCurrentNode->pChildNode[1] == NullNode)pChildNode = pCurrentNode->pChildNode[0];
            if (pChildNode != nullptr) {
                DeleteStack[++DeleteStackTop] = pCurrentNode;
                pChildNode->pFatherNode = pCurrentNode->pFatherNode;//如果c为nil，临时设置c的父节点
                return pCurrentNode = pChildNode;
            }
            pChildNode = pCurrentNode->pChildNode[1];
            while (pChildNode->pChildNode[0] != NullNode)pChildNode = pChildNode->pChildNode[0];
            pCurrentNode->key = pChildNode->key;
            pCurrentNode->mapped = pChildNode->mapped;
            DeleteStack[++DeleteStackTop] = pChildNode;
            pChildNode->pChildNode[1]->pFatherNode = pChildNode->pFatherNode;
            return pChildNode->pFatherNode->pChildNode[pChildNode->pFatherNode->pChildNode[0] !=
                                                       pChildNode] = pChildNode->pChildNode[1];
        }
        return _delete(pCurrentNode->pChildNode[pCurrentNode->key < key], key);
    }

    Mapped& _find(RBTree pCurrentNode,int key){
        if (pCurrentNode == NullNode)raise(1);
        if (pCurrentNode->key == key)return pCurrentNode->mapped;
        return _find(pCurrentNode->pChildNode[pCurrentNode->key < key], key);
    }

public:
    RedBlackTree() {
        DeleteStack = (RBTree *) malloc(128 * sizeof(RBTree));
        assert(DeleteStack);
        NullNode = (RBTree) malloc(sizeof(RedBlackTreeNode));
        assert(NullNode);
        DeleteStackTop = 0;
        NullNode->pChildNode[0] = NullNode->pChildNode[1] = NullNode;
        NullNode->pFatherNode = nullptr;
        NullNode->bNodeColor = BLACK;
        TreeRoot = NullNode;
    }

    ~RedBlackTree() {
        ClearDeleteStack();
        Clear();
        free(NullNode);
        free(DeleteStack);
    }

    void Insert(int key, int mapped) {
        RBTree pCurrentNode = _insert(TreeRoot, key, mapped, nullptr);
        if (pCurrentNode == nullptr)return;
        RBTree pFatherNode = pCurrentNode->pFatherNode, pGrandfatherNode;
        for (; pFatherNode != nullptr && pFatherNode->bNodeColor; pFatherNode = pCurrentNode->pFatherNode) {
            pGrandfatherNode = pFatherNode->pFatherNode;
            bool bRelationshipOfCurrentNodeAndFatherNode = pFatherNode->pChildNode[0] == pCurrentNode,
                    bRelationshipOfCurrentNodeAndGrandFatherNode = pGrandfatherNode->pChildNode[0] == pFatherNode;
            if (pGrandfatherNode->pChildNode[0]->bNodeColor == RED &&
                pGrandfatherNode->pChildNode[1]->bNodeColor == RED) {
                pGrandfatherNode->pChildNode[0]->bNodeColor = pGrandfatherNode->pChildNode[1]->bNodeColor = BLACK;
                pGrandfatherNode->bNodeColor = RED;
                pCurrentNode = pGrandfatherNode;
            } else {
                if (bRelationshipOfCurrentNodeAndFatherNode ^ bRelationshipOfCurrentNodeAndGrandFatherNode)
                    Rotate(pFatherNode, bRelationshipOfCurrentNodeAndFatherNode);
                pGrandfatherNode->bNodeColor = RED;
                pGrandfatherNode->pChildNode[!bRelationshipOfCurrentNodeAndGrandFatherNode]->bNodeColor = BLACK;
                Rotate(pGrandfatherNode, bRelationshipOfCurrentNodeAndGrandFatherNode);
                pCurrentNode = pGrandfatherNode->pChildNode[bRelationshipOfCurrentNodeAndGrandFatherNode];
            }
        }
        TreeRoot->bNodeColor = BLACK;
    }

    void Delete(int key) {
        RBTree pCurrentNode = _delete(TreeRoot, key);
        if (pCurrentNode == nullptr)return;
        if (DeleteStack[DeleteStackTop]->bNodeColor == BLACK) {
            RBTree pFatherNode, pBrotherNode;
            bool bRelationshipOfCurrentNodeAndFatherNode;
            while (pCurrentNode != TreeRoot && pCurrentNode->bNodeColor == BLACK) {
                pFatherNode = pCurrentNode->pFatherNode;
                bRelationshipOfCurrentNodeAndFatherNode = pFatherNode->pChildNode[0] != pCurrentNode;
                pBrotherNode = pFatherNode->pChildNode[!bRelationshipOfCurrentNodeAndFatherNode];
                if (pBrotherNode->bNodeColor) {
                    pBrotherNode->bNodeColor = BLACK;
                    pFatherNode->bNodeColor = RED;
                    Rotate(pFatherNode, bRelationshipOfCurrentNodeAndFatherNode);
                    pBrotherNode = pFatherNode->pChildNode[!bRelationshipOfCurrentNodeAndFatherNode];
                }
                if (pBrotherNode->pChildNode[0]->bNodeColor == BLACK &&
                    pBrotherNode->pChildNode[1]->bNodeColor == BLACK) {
                    pCurrentNode = pFatherNode;
                    pBrotherNode->bNodeColor = RED;
                } else {
                    if (pBrotherNode->pChildNode[!bRelationshipOfCurrentNodeAndFatherNode]->bNodeColor == BLACK) {
                        pBrotherNode->pChildNode[bRelationshipOfCurrentNodeAndFatherNode]->bNodeColor = BLACK;
                        pBrotherNode->bNodeColor = RED;
                        Rotate(pBrotherNode, !bRelationshipOfCurrentNodeAndFatherNode);
                        pBrotherNode = pFatherNode->pChildNode[!bRelationshipOfCurrentNodeAndFatherNode];
                    }
                    pBrotherNode->bNodeColor = pFatherNode->bNodeColor;
                    pFatherNode->bNodeColor = BLACK;
                    pBrotherNode->pChildNode[!bRelationshipOfCurrentNodeAndFatherNode]->bNodeColor = BLACK;
                    Rotate(pFatherNode, bRelationshipOfCurrentNodeAndFatherNode);
                    pCurrentNode = TreeRoot;
                }
            }
            pCurrentNode->bNodeColor = BLACK;
        }
        if (DeleteStackTop > 120)ClearDeleteStack();
    }

    void Clear(){
        if(TreeRoot==NullNode)return;
        _clear(TreeRoot);
        TreeRoot=NullNode;
    }

    void Print(RBTree pCurrentNode) {
        if (pCurrentNode == NullNode)return;
        Print(pCurrentNode->pChildNode[0]);
        printf("Node:%d Color:%d LeftChild:%d RightChild:%d FatherNode:%d Mapped:%d\n", pCurrentNode->key,
               pCurrentNode->bNodeColor, pCurrentNode->pChildNode[0]->key,
               pCurrentNode->pChildNode[1]->key, pCurrentNode->pFatherNode ? pCurrentNode->pFatherNode->key : 0,
               pCurrentNode->mapped);
        Print(pCurrentNode->pChildNode[1]);
    }

    Mapped Find(Key key){
        return _find(TreeRoot,key);
    }

    Mapped& operator[](Key key){
        return _find(TreeRoot,key);
    }
};

#ifdef DEBUG
struct ck {
    int data;
    bool color;
    int bklen;
    bool bnil;
} arr[10000000];

RedBlackTree<int, int>::RedBlackTreeNode *NullNode;

void _check(RedBlackTree<int, int>::RedBlackTreeNode *o, int d = 1) {
    if (o == NullNode) {
        arr[d].bklen = arr[d >> 1].bklen;
        arr[d].bnil = 1;
        return;
    }
    arr[d].data = o->key;
    arr[d].color = o->bNodeColor;
    arr[d].bklen = (!o->bNodeColor) + arr[d >> 1].bklen;
    if (o->pFatherNode && o->pFatherNode->bNodeColor && o->bNodeColor)puts("WA_1");
    _check(o->pChildNode[0], d << 1);
    _check(o->pChildNode[1], (d << 1) | 1);
}

void check(RedBlackTree<int, int> &t) {
    memset(arr, 0, sizeof arr);
    _check(t.TreeRoot);
    int last = -1;
    for (int i = 0; i < 10000000; ++i) {
        if (arr[i].bnil) {
            if (last != -1 && arr[i].bklen != last)puts("WA_2");
            last = arr[i].bklen;
        }
    }
}
#endif
int main() {
    RedBlackTree<int, int> rbt;
    rbt.Insert(233,555);
    printf("%d ",rbt[233]);
    rbt[233]=666;
    printf("%d ",rbt.Find(233));
    rbt.Delete(233);
#ifdef DEBUG
    NullNode = rbt.NullNode;
#endif
    for (int i = 0; i < 100000; ++i) {
        rbt.Insert(RandInt() % 233333, RandInt());
        rbt.Insert(RandInt() % 233333, RandInt());
        rbt.Insert(RandInt() % 233333, RandInt());
        rbt.Delete(RandInt() % 233333);
    }
#ifdef DEBUG
    puts("Checking");
    check(rbt);
#endif
    return 0;
}

[eiclpy]

树结构可视化

前提树节点左右子节点为C[0],C[1]。

#ifndef DATASTRUCT_PRINTTREE_H
#define DATASTRUCT_PRINTTREE_H

#include <cassert>
#include <cstdio>

template<class T>
class Printer {
    T *TreeRoot;
    T *NullNode;
    T **arrNodes;
    int MaxNodeDeep;

    void (*PrintNode)(T *node);

    int PrintWidth;

    void GetMaxNodeDeep(T *CurrentNode, int CurrentDeep = 0) {
        if (CurrentNode == NullNode) {
            if (CurrentDeep > MaxNodeDeep)MaxNodeDeep = CurrentDeep;
            return;
        }
        GetMaxNodeDeep(CurrentNode->c[0], CurrentDeep + 1);
        GetMaxNodeDeep(CurrentNode->c[1], CurrentDeep + 1);
    }

    void Init(T *CurrentNode, int Index = 1) {
        if (CurrentNode == NullNode)return;
        arrNodes[Index] = CurrentNode;
        Init(CurrentNode->c[0], Index << 1);
        Init(CurrentNode->c[1], Index << 1 | 1);
    }

    void PrintChars(int Count, char CharToPrint = ' ') {
        for (; Count-- > 0;)putchar(CharToPrint);
    }


public:
    Printer(T *root, T *nil, void (*printer)(T *node), int width) : TreeRoot(root), NullNode(nil),
                                                                        MaxNodeDeep(0), PrintNode(printer),
                                                                        PrintWidth(width + 2) {
        GetMaxNodeDeep(TreeRoot);
        arrNodes = new T *[1 << MaxNodeDeep];
        assert(arrNodes);
        for (int i = 0; i < 1 << MaxNodeDeep; ++i)arrNodes[i] = 0;
        Init(TreeRoot);
    }

    ~Printer() {
        assert(arrNodes);
        delete[]arrNodes;
    }

    void print() {
        for (int i = 1, d = 1; i <= MaxNodeDeep; ++i, d <<= 1) {
            int CurrentWidth = (1 << MaxNodeDeep - i - 1) * PrintWidth;
            int SpaceWidth = CurrentWidth;
            PrintChars(CurrentWidth >> 1);
            if (CurrentWidth == 0)CurrentWidth = 1;
            for (int j = d; j < d << 1; ++j) {
                putchar('|');
                PrintChars(CurrentWidth - PrintWidth >> 1, '-');
                arrNodes[j] == 0 ? PrintChars(PrintWidth - 2, ' ') : PrintNode(arrNodes[j]);
                PrintChars(CurrentWidth - PrintWidth >> 1, '-');
                putchar('|');
                PrintChars(SpaceWidth);
            }
            putchar('\n');
        }
    }
};

#endif DATASTRUCT_PRINTTREE_H