【C++ 学习 ㉔】- 详解 map 和 set（下）- map 和 set 的模拟实现

一、RBT.h

#pragma once
​
#include 
​
namespace yzz
{
    enum Color
    {
        RED,
        BLACK
    };
​
    template<class T>
    struct RBTNode
    {
        RBTNode* _left;
        RBTNode* _right;
        RBTNode* _parent;
        T _data;
        Color _clr;
​
        RBTNode(const T& data = T(), Color clr = RED)
            : _left(nullptr), _right(nullptr), _parent(nullptr)
            , _data(data), _clr(clr)
        { }
    };
​
    template<class T, class Ref, class Ptr>
    struct RBTIterator
    {
        typedef RBTNode Node;
        typedef RBTIterator iterator;
        typedef RBTIteratorconst T&, const T*> const_iterator;
        typedef RBTIterator self;
​
        Node* _pnode;
​
        RBTIterator(Node* p) : _pnode(p) { }
​
        // 可以通过 iterator 迭代器构造 const_iterator 迭代器
        RBTIterator(const iterator& it) : _pnode(it._pnode) { }
​
        Ref operator*() const
        {
            return _pnode->_data;
        }
​
        T* operator->() const
        {
            return &_pnode->_data;
        }
​
        void increment()
        {
            if (_pnode->_right)
            {
                // 若右子树存在，则找到右子树中值最小的节点（即右子树中最左侧的节点）
                Node* rightMin = _pnode->_right;
                while (rightMin->_left)
                {
                    rightMin = rightMin->_left;
                }
                _pnode = rightMin;
            }
            else
            {
                // 若右子树不存在，表明以 *_pnode 为根的子树已经访问完了
                Node* parent = _pnode->_parent;
                while (parent)
                {
                    if (_pnode == parent->_left)
                    {
                        // 若 *_pnode 是 *parent 的左孩子，
                        // 则下一个该访问的节点就是 *parent
                        break;
                    }
                    else
                    {
                        // 若 *_pnode 是 *parent 的右孩子，
                        // 则以 *parent 为根的子树也已经访问了
                        _pnode = parent;
                        parent = parent->_parent;
                    }
                }
                _pnode = parent;
            }
        }
​
        void decrement()
        {
            if (_pnode->_left)
            {
                // 若左子树存在，则找到左子树中值最大的节点（即左子树中最右侧的节点）
                Node* leftMax = _pnode->_left;
                while (leftMax->_right)
                {
                    leftMax = leftMax->_right;
                }
                _pnode = leftMax;
            }
            else
            {
                Node* parent = _pnode->_parent;
                while (parent)
                {
                    if (_pnode == parent->_right)
                    {
                        // 若 *_pnode 是 *parent 的右孩子，
                        // 则上一个已访问的节点就是 *parent
                        break;
                    }
                    else
                    {
                        // 若 *_pnode 是 *parent 的左孩子
                        _pnode = parent;
                        parent = parent->_parent;
                    }
                }
                _pnode = parent;
            }
        }
​
        self& operator++()
        {
            increment();
            return *this;
        }
​
        self operator++(int)
        {
            self tmp = *this;
            increment();
            return tmp;
        }
​
        self& operator--()
        {
            decrement();
            return *this;
        }
​
        self operator--(int)
        {
            self tmp = *this;
            decrement();
            return tmp;
        }
​
        bool operator==(const self& it) const
        {
            return _pnode == it._pnode;
        }
​
        bool operator!=(const self& it) const
        {
            return _pnode != it._pnode;
        }
    };
​
    template<class K, class T, class KOfT>
    class RBT
    {
        typedef RBTNode Node;
    public:
        /*---------- 构造函数和析构函数 ----------*/
        RBT() : _root(nullptr) { }
​
        ~RBT()
        {
            Destroy(_root);
        }
​
        /*---------- 迭代器 ----------*/
        typedef RBTIterator iterator;
        typedef RBTIteratorconst T&, const T*> const_iterator;
​
        iterator begin()
        {
            Node* leftMin = _root;
            while (leftMin->_left)
            {
                leftMin = leftMin->_left;
            }
            return iterator(leftMin);
        }
​
        iterator end()
        {
            return iterator(nullptr);
        }
​
        const_iterator begin() const
        {
            Node* leftMin = _root;
            while (leftMin->_left)
            {
                leftMin = leftMin->_left;
            }
            return const_iterator(leftMin);
        }
​
        const_iterator end() const
        {
            return const_iterator(nullptr);
        }
​
        /*---------- 插入 ----------*/
        std::pairbool> insert(const T& data)
        {
            if (_root == nullptr)
            {
                _root = new Node(data, BLACK);
                return std::make_pair(iterator(_root), true);
            }
​
            Node* parent = nullptr;
            Node* cur = _root;
            KOfT kot;
            while (cur)
            {
                parent = cur;
                if (kot(data) < kot(cur->_data))
                    cur = cur->_left;
                else if (kot(data) > kot(cur->_data))
                    cur = cur->_right;
                else
                    return std::make_pair(iterator(cur), false);
            }
​
            // 如果插入前树非空，新插入的节点应该是红色节点
            cur = new Node(data, RED);
            Node* tmp = cur;
            if (kot(data) < kot(parent->_data))
                parent->_left = cur;
            else
                parent->_right = cur;
            cur->_parent = parent;
​
            // 出现连续两个红色节点的情形
            while (parent && parent->_clr == RED)
            {
                Node* grandparent = parent->_parent;
                Node* uncle;
                if (grandparent->_left == parent)
                    uncle = grandparent->_right;
                else
                    uncle = grandparent->_left;
​
                // 如果 *uncle 是红色节点
                if (uncle && uncle->_clr == RED)
                {
                    parent->_clr = uncle->_clr = BLACK;
                    grandparent->_clr = RED;
​
                    cur = grandparent;
                    parent = cur->_parent;
                }
                else  // 如果 *uncle 不存在或者是黑色节点
                {
                    if (grandparent->_left == parent && parent->_left == cur)
                    {
                        LL(grandparent);
                        parent->_clr = BLACK;
                        grandparent->_clr = RED;
                    }
                    else if (grandparent->_right == parent && parent->_right == cur)
                    {
                        RR(grandparent);
                        parent->_clr = BLACK;
                        grandparent->_clr = RED;
                    }
                    else if (grandparent->_left == parent && parent->_right == cur)
                    {
                        LR(grandparent);
                        cur->_clr = BLACK;
                        grandparent->_clr = RED;
                    }
                    else if (grandparent->_right == parent && parent->_left == cur)
                    {
                        RL(grandparent);
                        cur->_clr = BLACK;
                        grandparent->_clr = RED;
                    }
                    break;
                }
            }
            _root->_clr = BLACK;
            return std::make_pair(iterator(tmp), true);
        }
        
    private:
        void Destroy(Node*& root)
        {
            // 【注意:root 为 _root 或者某个节点的左或右指针的引用】
            if (root == nullptr)
                return;
​
            Destroy(root->_left);
            Destroy(root->_right);
            delete root;
            root = nullptr;
        }
​
        void LL(Node* parent)
        {
            Node* cur = parent->_left;
            Node* curRight = cur->_right;
​
            parent->_left = curRight;
            if (curRight)
                curRight->_parent = parent;
​
            cur->_right = parent;
            Node* tmp = parent->_parent;
            parent->_parent = cur;
            cur->_parent = tmp;
​
            if (tmp == nullptr)
            {
                _root = cur;
            }
            else
            {
                if (tmp->_left == parent)
                    tmp->_left = cur;
                else
                    tmp->_right = cur;
            }
        }
​
        void RR(Node* parent)
        {
            Node* cur = parent->_right;
            Node* curLeft = cur->_left;
​
            parent->_right = curLeft;
            if (curLeft)
                curLeft->_parent = parent;
​
            cur->_left = parent;
            Node* tmp = parent->_parent;
            parent->_parent = cur;
            cur->_parent = tmp;
​
            if (tmp == nullptr)
            {
                _root = cur;
            }
            else
            {
                if (tmp->_left == parent)
                    tmp->_left = cur;
                else
                    tmp->_right = cur;
            }
        }
​
        void LR(Node* parent)
        {
            Node* cur = parent->_left;
            RR(cur);
            LL(parent);
        }
​
        void RL(Node* parent)
        {
            Node* cur = parent->_right;
            LL(cur);
            RR(parent);
        }
    private:
        Node* _root;
    };
}

二、set.h

#pragma once
​
#include "RBT.h"
​
namespace yzz
{
    template<class K>
    class set
    {
        struct SetKOfT
        {
            const K& operator()(const K& key)
            {
                return key;
            }
        };
        typedef RBT RBT;
    public:
        typedef typename RBT::const_iterator iterator;
        typedef typename RBT::const_iterator const_iterator;
​
        const_iterator begin() const
        {
            return _t.begin();
        }
​
        const_iterator end() const
        {
            return _t.end();
        }
​
        std::pairbool> insert(const K& key)
        {
            std::pair<typename RBT::iterator, bool> ret = _t.insert(key);
            return std::pairbool>(ret.first, ret.second);
        }
    private:
        RBT _t;
    };
}

三、map.h

#pragma once
​
#include "RBT.h"
​
namespace yzz
{
    template<class K, class V>
    class map
    {
        struct MapKOfT
        {
            const K& operator()(const std::pair<const K, V>& kv)
            {
                return kv.first;
            }
        };
        typedef RBTconst K, V>, MapKOfT> RBT;
    public:
        typedef typename RBT::iterator iterator;
        typedef typename RBT::const_iterator const_iterator;
​
        iterator begin()
        {
            return _t.begin();
        }
​
        iterator end()
        {
            return _t.end();
        }
​
        const_iterator begin() const
        {
            return _t.begin();
        }
​
        const_iterator end() const
        {
            return _t.end();
        }
​
        std::pairbool> insert(const std::pair<const K, V>& kv)
        {
            return _t.insert(kv);
        }
​
        V& operator[](const K& key)
        {
            std::pairbool> ret = _t.insert(std::make_pair(key, V()));
            return ret.first->second;
        }
    private:
        RBT _t;
    };
}

四、test.cpp

#include "set.h"
#include "map.h"
#include 
using namespace std;
​
void TestSet()
{
    int arr[] = { 16, 3, 7, 11, 9, 26, 18, 14, 15 };
    yzz::set<int> s;
    for (const auto& e : arr)
    {
        s.insert(e);
    }
    yzz::set<int>::iterator it = s.begin();
    while (it != s.end())
    {
        // *it = 0;  // error
        cout << *it << " ";
        ++it;
    }
    // 3 7 9 11 14 15 16 18 26
    cout << endl;
}
​
void PrintMap(const yzz::map<int, int>& m)
{
    yzz::map<int, int>::const_iterator it = m.begin();
    while (it != m.end())
    {
        // it->first = 1;  // error
        // it->second = 2;  // error
        cout << it->first << " : " << it->second << endl;
        ++it;
    }
}
​
void TestMap()
{
    int arr[] = { 16, 3, 7, 11, 9, 26, 18, 14, 15 };
    yzz::map<int, int> m;
    for (const auto& e : arr)
    {
        m.insert(make_pair(e, e));
    }
    yzz::map<int, int>::iterator it = m.begin();
    while (it != m.end())
    {
        // it->first = 1;  // error
        it->second = 2;  // ok
        cout << it->first << " : " << it->second << endl;
        ++it;
    }
​
    for (const auto& kv : m)
    {
        m[kv.first] = kv.first;
    }
    PrintMap(m);
}
​
int main()
{
    TestSet();
    TestMap();
    return 0;
}