C++精通之路：模拟实现map\u002Fset

//颜色
enum Colour
{
    RED,
    BLACK,
};

template<class T>
struct RBTreeNode
{
    RBTreeNode<T>* _left;
    RBTreeNode<T>* _right;
    RBTreeNode<T>* _parent;

    T _data;//T可以是key也可以是pair<K,V>
    Colour _col;

    RBTreeNode(const T&amp; data)
        :_left(nullptr)
        , _right(nullptr)
        , _parent(nullptr)
        , _data(data)
        , _col(RED)
    {}
};

template<class K, class T, class KeyOfT>
class RBTree
{
    typedef RBTreeNode<T> Node;
public:
    typedef _TreeIterator<T, T&amp;, T*> iterator;
    typedef _TreeIterator<T,const T&amp;, const T*> const_iterator;
    typedef ReverseIterator<iterator> reverse_iterator;
    typedef ReverseIterator<const_iterator> reverse_const_iterator;

    RBTree()
        :_root(nullptr)
    {}

    ~RBTree()
    {
        _Destory(_root);
    }

    iterator begin()
    {
        Node* cur = _root;
        while (cur&amp;&amp;cur->_left)
        {
            cur = cur->_left;
        }
        return iterator(cur);
    }

    reverse_iterator rbegin()
    {
        Node* cur = _root;
        while (cur&amp;&amp;cur->_right)
        {
            cur = cur->_right;
        }
        return reverse_iterator(iterator(cur));
    }

    reverse_iterator rend()
    {
        return reverse_iterator(iterator(nullptr));
    }

    iterator end()
    {
        return iterator(nullptr);
    }

    Node* Find(const K&amp; key)
    {
        KeyOfT kot;
        Node* cur = _root;
        while (cur)
        {
            if (kot(cur->_kv.first) > key)
            {
                cur = cur->_left;
            }
            else if (kot(cur->_kv.first) < key)
            {
                cur = cur->_right;
            }
            else
            {
                return cur;
            }
        }
        return nullptr;
    }

    pair<iterator, bool> Insert(const T&amp; data)
    {
        //空树的情况
        if (_root == nullptr)
        {
            _root = new Node(data);
            _root->_col = BLACK;
            return make_pair(iterator(_root), true);
        }
        KeyOfT kot;
        //查找位置插入节点
        Node* cur = _root, * parent = _root;
        while (cur)
        {
            if (kot(cur->_data) > kot(data))
            {
                parent = cur;
                cur = cur->_left;
            }
            else if (kot(cur->_data) < kot(data))
            {
                parent = cur;
                cur = cur->_right;
            }
            else
            {
                return make_pair(iterator(cur), false);
            }
        }

        //创建链接节点
        cur = new Node(data);
        Node* newnode = cur;
        if (kot(parent->_data) > kot(data))
        {
            parent->_left = cur;
        }
        else
        {
            parent->_right = cur;
        }
        cur->_parent = parent;

        //父节点存在且为红，则需要调整（不能存在连续的红色节点）
        while (parent &amp;&amp; parent->_col == RED)
        {
            //此时当前节点一定有祖父节点
            Node* granparent = parent->_parent;
            //具体调整情况主要看叔叔节点
            //分左右讨论
            if (parent == granparent->_left)
            {
                Node* uncle = granparent->_right;
                //情况1：叔叔节点存在且为红
                if (uncle &amp;&amp; uncle->_col == RED)
                {
                    //修改颜色，继续向上检查
                    granparent->_col = RED;
                    parent->_col = uncle->_col = BLACK;

                    cur = granparent;
                    parent = cur->_parent;
                }
                else//情况2和3：叔叔节点不存在 或者存在且为黑
                {
                    //单旋(三代节点为斜线)+变色
                    if (cur == parent->_left)
                    {
                        RotateR(granparent);

                        granparent->_col = RED;
                        parent->_col = BLACK;
                    }
                    else//双旋(三代节点为折线)+变色
                    {
                        RotateL(parent);
                        RotateR(granparent);

                        cur->_col = BLACK;
                        granparent->_col = RED;
                    }
                    //旋转后不需再向上调整了
                    break;
                }
            }
            else//parent=grandparent->right
            {
                Node* uncle = granparent->_left;
                if (uncle &amp;&amp; uncle->_col == RED)
                {
                    parent->_col = uncle->_col = BLACK;
                    granparent->_col = RED;

                    cur = granparent;
                    parent = cur->_parent;
                }
                else
                {
                    if (cur == parent->_right)
                    {
                        RotateL(granparent);

                        parent->_col = BLACK;
                        granparent->_col = RED;
                    }
                    else
                    {
                        RotateR(parent);
                        RotateL(granparent);

                        cur->_col = BLACK;
                        granparent->_col = RED;
                    }
                    break;
                }
            }
        }

        //确保根节点为黑
        _root->_col = BLACK;
        return make_pair(iterator(newnode), true);
    }

    bool IsRBTree()
    {
        if (_root == nullptr)
        {
            return true;
        }

        if (_root->_col == RED)
        {
            cout << "根节点为红色" << endl;
            return false;
        }

        int Blacknum = 0;
        Node* cur = _root;
        while (cur)
        {
            if (cur->_col == BLACK)
                Blacknum++;
            cur = cur->_left;
        }

        int i = 0;
        return _IsRBTree(_root, Blacknum, i);
    }

private:
    void _Destory(Node*&amp; root)
    {
        if (root == nullptr)
            return;

        _Destory(root->_left);
        _Destory(root->_right);

        delete root;
        root = nullptr;
    }
    
    bool _IsRBTree(Node* root, int blacknum, int count)
    {
        if (root == nullptr)
        {
            if (blacknum == count)
                return true;
            cout << "各路径上黑色节点个数不同" << endl;
            return false;
        }

        if (root->_col == RED &amp;&amp; root->_parent->_col == RED)
        {
            cout << "存在连续红色节点" << endl;
            return false;
        }

        if (root->_col == BLACK)
            count++;

        return _IsRBTree(root->_left, blacknum, count) &amp;&amp; _IsRBTree(root->_right, blacknum, count);
    }

    void RotateL(Node* parent)
    {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        Node* parentP = parent->_parent;


        parent->_right = subRL;
        if (subRL)
        {
            subRL->_parent = parent;
        }

        subR->_left = parent;
        parent->_parent = subR;

        if (parent == _root)
        {
            _root = subR;
            subR->_parent = nullptr;
        }
        else
        {
            subR->_parent = parentP;
            if (parentP->_left == parent)
            {
                parentP->_left = subR;
            }
            else
            {
                parentP->_right = subR;
            }
        }
    }

    void RotateR(Node* parent)
    {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        Node* parentP = parent->_parent;


        parent->_left = subLR;
        if (subLR)
        {
            subLR->_parent = parent;
        }

        subL->_right = parent;
        parent->_parent = subL;

        if (parent == _root)
        {
            _root = subL;
            subL->_parent = nullptr;
        }
        else
        {
            subL->_parent = parentP;
            if (parentP->_left == parent)
            {
                parentP->_left = subL;
            }
            else
            {
                parentP->_right = subL;
            }
        }
    }

private:
    Node* _root;
};

namespace ymhh
{
    template<class K, class V>
    class map
    {
        struct MapOfKey
        {
            const K&amp; operator()(const pair<K, V>&amp; kv)
            {
                return kv.first;
            }
        };
    public:
        typedef typename RBTree<K, pair<const K, V>, MapOfKey>::iterator iterator;
        typedef typename RBTree<K, pair<const K, V>, MapOfKey>::reverse_iterator reverse_iterator;

        iterator begin()
        {
            return _t.begin();
        }

        iterator end()
        {
            return _t.end();
        }

        reverse_iterator rbegin()
        {
            return _t.rbegin();
        }

        reverse_iterator rend()
        {
            return _t.rend();
        }
        pair<iterator, bool> insert(const pair<const K, V>&amp; kv)
        {
            return _t.Insert(kv);
        }

        V&amp; operator[](const K&amp; key)
        {
            pair<iterator, bool> ret = insert(make_pair(key, V()));
            return ret.first->second;
        }

        iterator find(const K&amp; key)
        {
            return _t.Find(key);
        }

    private:
        RBTree<K, pair<const K, V>, MapOfKey> _t;
    };
}

namespace ymhh
{
    template<class K>
    class set
    {
        struct SetOfKey
        {
            const K&amp; operator()(const K&amp; key)
            {
                return key;
            }
        };
    public:
        typedef typename RBTree<K,K, SetOfKey>::iterator iterator;
        typedef typename RBTree<K,K, SetOfKey>::reverse_iterator reverse_iterator;

        iterator begin()
        {
            return _t.begin();
        }

        iterator end()
        {
            return _t.end();
        }

        reverse_iterator rbegin()
        {
            return _t.rbegin();
        }

        reverse_iterator rend()
        {
            return _t.rend();
        }

        pair<iterator, bool> insert(const K&amp; key)
        {
            return _t.Insert(key);
        }

        iterator find(const K&amp; key)
        {
            return _t.Find(key);
        }

    private:
        RBTree<K, K, SetOfKey> _t;
    };
}