用红黑树实现map和set

用红黑树实现map和set

上一篇文章讲解了如何实现红黑树：
https://blog.csdn.net/Radium_1209/article/details/104873813
这里我们用已经实现的红黑树来写一个简单的map和set。
因为map有两个参数，所以我们要先对原来的代码进行微调，将传入的参数调整为Key和Value。
还有一些地方需要微调，详见https://github.com/Radium1209/Red-Black-Tree

红黑树的接口定义

#ifndef RED_BLACK_TREE
#define RED_BALCK_TREE

#include <stdlib.h>

/* 红黑树颜色 */
enum RBTColor
{
    RED,
    BLACK,
};

/* 红黑树节点 */
template<class Key, class Value>
class RBTNode
{
public:
    RBTNode( RBTColor col, Key k, Value val, RBTNode* lch, RBTNode* rch, RBTNode* fa )
    : color( col ), key( k ), value( val ), blackNum( 0 ), leftChild( lch ), rightChild( rch ), father( fa ) {}

    RBTColor color;         // 颜色(红或黑)
    Key key;                  // 关键字
    Value value;                // 数据
    int blackNum;           // 该节点的路径包含的黑色节点个数
    RBTNode *leftChild;     // 左孩子
    RBTNode *rightChild;    // 右孩子
    RBTNode *father;        // 父节点
};

/* 节点-层数对，用于遍历 */
template<class Key, class Value>
class NLPair
{
public:
    NLPair( RBTNode<Key, Value>* node, int layer )
    : node( node ), layer( layer ) {}
	
    RBTNode<Key, Value>* node;
    unsigned int layer;
};

template<class Key, class Value>
class RBTree
{
    
private:
    RBTNode<Key, Value> *root;

public:
    RBTree();
    ~RBTree();
    // 判断是否为红黑树
    bool isRBT();
    // 各种插入函数
    void insert( Key key );
    void insert( Key key, Value value );
    void insertUnique( Key key );
    void insertUnique( Key key, Value value );
    // 删除函数
    void erase( Key key );
    // 清空
    void clear();
    // 查询函数
    RBTNode<Key, Value>* find( Key key );
    // 画红黑树，测试用
    void draw();
    // 求开始结束节点，此处没有实现iterator，所以用处不大
    RBTNode<Key, Value>* begin();
    RBTNode<Key, Value>* end();
    // 求节点个数和是否为空
    unsigned int size();
    bool empty();

// 一些内部函数
private:
    bool isBST( RBTNode<Key, Value> *node );
    bool hasTwoRed( RBTNode<Key, Value> *node );
    bool hasSameBlack( RBTNode<Key, Value> *node );
    void calBlackNum( RBTNode<Key, Value> *node, int blackNum );
    void leftRotate( RBTNode<Key, Value> *root );
    void rightRotate( RBTNode<Key, Value> *root );
    void insert( RBTNode<Key, Value> *node );
    void insertFix( RBTNode<Key, Value> *node );
    void erase( RBTNode<Key, Value> *node );
    void eraseFix( RBTNode<Key, Value> *node );
    void eraseNode( RBTNode<Key, Value> *node );
    void clear( RBTNode<Key, Value> *node );
    RBTNode<Key, Value>* find( RBTNode<Key, Value> *node, Key key );
    unsigned int size( RBTNode<Key, Value> *node );
};
#include "rbTree.cpp"
#endif

set

set完全直接调用红黑树的接口即可，不用进行任何修改。我们只实现了简单的插入删除查询等函数。因为set是没有重复元素的，我们可以模仿STL实现两个insert函数，一个可以插入重复的key，另一个只能插入不重复的节点。

// mySet.h

#ifndef MY_SET
#define MY_SET

template<class Key>
class mySet
{
private:
    RBTree<Key, Key>* rbt;

public:
    mySet();
    ~mySet();

public:
    RBTNode<Key, Key>* begin();
    RBTNode<Key, Key>* end();
    void clear();
    void insert( Key key );
    void erase( Key key );
    unsigned int size();
    RBTNode<Key, Key>* find( Key key );
    bool empty();
};
#include "mySet.cpp"
#endif

template<class Key>
mySet<Key>::mySet()
{
    rbt = new RBTree<Key, Key>();
}

template<class Key>
mySet<Key>::~mySet()
{
    delete rbt;
}

template<class Key>
RBTNode<Key, Key>* mySet<Key>::begin()
{
    return rbt->begin();
}

template<class Key>
RBTNode<Key, Key>* mySet<Key>::end()
{
    return rbt->end();
}

template<class Key>
void mySet<Key>::clear()
{
    rbt->clear();
}

template<class Key>
void mySet<Key>::insert( Key key )
{
    rbt->insertUnique( key );
}

template<class Key>
void mySet<Key>::erase( Key key )
{
    rbt->erase( key );
}

template<class Key>
unsigned int mySet<Key>::size()
{
    return rbt->size();
}

template<class Key>
RBTNode<Key, Key>* mySet<Key>::find( Key key )
{
    return rbt->find( key );
}

template<class Key>
bool mySet<Key>::empty()
{
    return rbt->empty();
}

map

map的话也基本上差不多，也是直接用红黑树就可以。这边重载了一下[]，使其可以直接赋值。修改了一下之前的insert函数，使其返回一个指针。

// myMap.h

#ifndef MY_MAP
#define MY_MAP

template<class Key, class Value>
class myMap
{
private:
    RBTree<Key, Value>* rbt;

public:
    myMap();
    ~myMap();

public:
    RBTNode<Key, Value>* begin();
    RBTNode<Key, Value>* end();
    void clear();
    void insert( Key key, Value value );
    void erase( Key key );
    unsigned int size();
    RBTNode<Key, Value>* find( Key key );
    bool empty();
    Value& operator[] ( Key key );
};
#include "myMap.cpp"
#endif

// myMap.cpp

template<class Key, class Value>
myMap<Key, Value>::myMap()
{
    rbt = new RBTree<Key, Value>();
}

template<class Key, class Value>
myMap<Key, Value>::~myMap()
{
    delete rbt;
}

template<class Key, class Value>
RBTNode<Key, Value>* myMap<Key, Value>::begin()
{
    return rbt->begin();
}

template<class Key, class Value>
RBTNode<Key, Value>* myMap<Key, Value>::end()
{
    return rbt->end();
}

template<class Key, class Value>
void myMap<Key, Value>::clear()
{
    rbt->clear();
}

template<class Key, class Value>
void myMap<Key, Value>::insert( Key key, Value value )
{
    rbt->insert( key, value );
}

template<class Key, class Value>
void myMap<Key, Value>::erase( Key key )
{
    rbt->erase( key );
}

template<class Key, class Value>
unsigned int myMap<Key, Value>::size()
{
    return rbt->size();
}

template<class Key, class Value>
RBTNode<Key, Value>* myMap<Key, Value>::find( Key key )
{
    return find( key );
}

template<class Key, class Value>
bool myMap<Key, Value>::empty()
{
    return rbt->empty();
}

template<class Key, class Value>
Value& myMap<Key, Value>::operator[] ( Key key )
{
    return rbt->insertUnique( key )->value;
}

测试

这里还添加了一个简单的测试

#include "rbTree.h"
#include "mySet.h"
#include "myMap.h"
#include <iostream>
#include <string>
using namespace std;

int main( int argc, char* argv[] )
{
    cout << "set test" << endl;
    mySet<int> s;
    s.insert(1);
    cout << (s.find(1))->key << endl;
    s.insert(10);
    cout << (s.end())->key << endl;
    s.insert(23);
    cout << (s.begin())->key << endl;
    s.erase(23);
    cout << (s.end())->key << endl;

    cout << endl << "map test" << endl;
    myMap<int, int> m;
    m.insert(1, 100);
    m[5] = 132;
    m.insert(3, 11);
    m.insert(11, 44);
    cout << "m[5]=" << m[5] << endl;
    cout << "m[3]=" << m[3] << endl;
    m.erase(5);
    m.erase(100);

    cout << endl << "rbTree test" << endl;
    RBTree<int, int> tree;
    string option;
    for (int i=1; i<=10000; i++)
    {
        tree.insert( i );
    }
    for (int i=1; i<=10000; i++)
    {
        tree.erase( i );
    }
    while( true )
    {
        cout << "> ";
        cin >> option;
        if ( option == "insert" )
        {
            int k, val;
            cin >> k >> val;
            tree.insert( k, val );
            cout << "finish insert" << endl;
            tree.isRBT();
        }
        else if ( option == "erase" )
        {
            int val;
            cin >> val;
            tree.erase( val );
            cout << "finish erase" << endl;
            tree.isRBT();
        }
        else if ( option == "draw" )
        {
            tree.draw();
        }
        else if ( option == "exit" )
        {
            break;
        }
        
    }
    return 0;
}