代码如下:
#pragma once
#include <iostream>
using namespace std;enum COLOR
{BLACK, RED
};template<class V>//迭代器声明,定义在后面
struct RBTreeIterator;template<typename V>
struct RBTreeNode
{RBTreeNode<V> * _parent;RBTreeNode<V>* _left;RBTreeNode<V> *_right;V _val;COLOR _color;RBTreeNode(const V &val = V()) :_parent(nullptr), _left(nullptr), _right(nullptr), _val(val), _color(RED) {}};template<typename K, typename V, typename KeyOfValue>
class RBTree
{
public:typedef RBTreeNode<V> Node;typedef RBTreeIterator<V> iterator;RBTree() :_header(new Node){_header->_left = _header->_right = _header;}pair<iterator, bool> insert(const V &val){if (_header->_parent == nullptr){Node * root = new Node(val);_header->_parent = root;root->_parent = _header;_header->_left = _header->_right = root;//根结点为黑色root->_color = BLACK;return make_pair(iterator(root), true);}Node *cur = _header->_parent;Node *parent = nullptr;KeyOfValue kov;//1.寻找要插入的结点的位置while (cur){parent = cur;if (kov(cur->_val) == kov(val)) return make_pair(iterator(cur), false);else if (kov(cur->_val) > kov(val)) cur = cur->_left;else cur = cur->_right;}//2.创建结点cur = new Node(val);Node *node = cur;//记录插入的结点的位置,方便后面返回。if (kov(parent->_val) > kov(cur->_val)) parent->_left = cur;else parent->_right = cur;cur->_parent = parent;//3.颜色的修改或者结构的调整while (cur != _header->_parent && cur->_parent->_color == RED){parent = cur->_parent;Node *gfather = parent->_parent;if (gfather->_left == parent){Node *uncle = gfather->_right;//情况1.uncle存在且为红if (uncle && uncle->_color == RED){parent->_color = uncle->_color = BLACK;gfather->_color = RED;//向上追溯cur = gfather;}else{if (parent->_right == cur)//情况3{RotateL(parent);swap(cur, parent);}//2.uncle不存在或者uncle为黑RotateR(gfather);parent->_color = BLACK;gfather->_color = RED;break;}}else{Node *uncle = gfather->_left;if (uncle && uncle->_color == RED){parent->_color = uncle->_color = BLACK;gfather->_color = RED;cur = gfather;}else{if (parent->_left == cur){RotateR(parent);swap(cur, parent);}RotateL(gfather);parent->_color = BLACK;gfather->_color = RED;break;}}}//根结点为黑色_header->_parent->_color = BLACK;//更新头结点的左右指向_header->_left = leftMost();_header->_right = rightMost();return make_pair(iterator(node), true);}void RotateL(Node *parent){Node *subR = parent->_right;Node *subRL = subR->_left;parent->_right = subRL;if (subRL) subRL->_parent = parent;if (parent == _header->_parent){_header->_parent = subR;subR->_parent = _header;}else{Node *gfather = parent->_parent;if (gfather->_left == parent) gfather->_left = subR;else gfather->_right = subR;subR->_parent = gfather;}subR->_left = parent;parent->_parent = subR;}void RotateR(Node * parent){Node *subL = parent->_left;Node *subLR = subL->_right;parent->_left = subLR;if (subLR) subLR->_parent = parent;if (parent == _header->_parent){_header->_parent = subL;subL->_parent = _header;}else{Node *gfather = parent->_parent;if (gfather->_left == parent) gfather->_left = subL;else gfather->_right = subL;subL->_parent = gfather;}subL->_right = parent;parent->_parent = subL;}Node *leftMost(){Node *cur = _header->_parent;while (cur && cur->_left){cur = cur->_left;}return cur;}Node *rightMost(){Node *cur = _header->_parent;while (cur && cur->_right){cur = cur->_right;}return cur;}iterator begin(){return iterator(_header->_left);}iterator end(){return iterator(_header);}//反向迭代器iterator rbegin(){return iterator(_header->_right);}private:Node *_header;
};template<class V>
struct RBTreeIterator
{typedef RBTreeNode<V> Node;typedef RBTreeIterator<V> Self;Node *_node;RBTreeIterator(Node *node) :_node(node) {}V & operator*(){return _node->_val;}V * operator->(){return &_node->_val;}bool operator!=(const Self & it){return _node != it._node;}Self & operator++(){if (_node->_right)//存在右结点{//右子树的最左结点_node = _node->_right;while (_node->_left){_node = _node->_left;}}else//不存在右结点{Node *parent = _node->_parent;while (_node == parent->_right)//回溯{_node = parent;parent = parent->_parent;}//特殊情况:根结点没有右孩子,则不需要更新结点if (_node->_right != parent) _node = parent;}return *this;}Self & operator--(){if (_node->_left){//右子树的最左结点_node = _node->_left;while (_node->_right){_node = _node->_right;}}else{Node *parent = _node->_parent;while (_node == parent->_left){_node = parent;parent = parent->_parent;}if (_node->_left != parent) _node = parent;}return *this;}
};template<typename K, typename T>
class Map
{struct MapKeyOfValue{const K& operator()(const pair<K, T> & val){return val.first;}};public:typedef typename RBTree<K, pair<K, T>, MapKeyOfValue>::iterator iterator;pair<iterator, bool> insert(const pair<K, T> & kv){return _rbt.insert(kv);}T & operator[](const K &key){pair<iterator, bool> ret = _rbt.insert(make_pair(key, T()));return ret.first->second;}iterator begin(){return _rbt.begin();}iterator end(){return _rbt.end();}iterator rbegin(){return _rbt.rbegin();}private:typedef RBTree<K, pair<K, T>, MapKeyOfValue> rbt;rbt _rbt;
};template<typename K>
class Set
{struct SetKeyOfValue{const K & operator()(const K & val){return val;}};public:bool insert(const K &val){return _rbt.insert(val);}private:typedef RBTree<K, K, SetKeyOfValue> rbt;rbt _rbt;};
测试代码如下:
#include<iostream>
#include "Map.h"
using namespace std;int main()
{Map<int, int> m;m.insert(make_pair(1, 1));m.insert(make_pair(2, 1));m.insert(make_pair(3, 1));m.insert(make_pair(4, 1));Map<int, int>::iterator it = m.begin();while (it != m.end()){//it-> 获取结点的val的地址 => pair指针//pair->first获得pair对象的第一个值//pair->second获得pair对象的第二个值cout << it->first << " " << it->second << endl;++it;}Map<int, int>::iterator it1 = m.rbegin();while (it1 != m.end()){cout << it1->first << " " << it1->second << endl;--it1;}Map<int, int> m1;m1[1] = 1;m1[3] = 20;m1[2] = 300;m1[4] = 4000;Map<int, int>::iterator it2 = m1.begin();while (it2 != m1.end()){cout << it2->first << " " << it2->second << endl;++it2;}return 0;}
测试结果: