从C语言到C++_29(红黑树封装set和map)红黑树迭代器的实现（下）

#pragma once
 
#include <iostream>
#include <assert.h>
#include <time.h>
using namespace std;
 
enum Colour // 枚举颜色
{
  RED,
  BLACK
};
 
template<class T>
struct RBTreeNode
{
public:
  RBTreeNode<T>* _left;
  RBTreeNode<T>* _right;
  RBTreeNode<T>* _parent;
 
  T _data; // 结点中的数据
  Colour _col; // 比AVL树少了平衡因子，多了颜色
 
  RBTreeNode(const T& data)
    :_left(nullptr)
    , _right(nullptr)
    , _parent(nullptr)
    , _data(data)
  {}
};
 
template<class T, class Ref, class Ptr>
struct __RBTreeIterator
{
public:
  typedef RBTreeNode<T> Node;
  typedef __RBTreeIterator<T, Ref, Ptr> Self;
  Node* _node;
 
  __RBTreeIterator(Node* node)
    :_node(node)
  {}
 
  Ref operator*()
  {
    return _node->_data;
  }
 
  Ptr operator->()
  {
    return &_node->_data;
  }
 
  bool operator!=(const Self& s) const
  {
    return _node != s._node;
  }
 
  bool operator==(const Self& s) const
  {
    return _node == s._node;
  }
 
  Self& operator++()
  {
    if (_node->_right) // 右子树存在，++后就到右子树的最左结点
    {
      Node* left = _node->_right;
      while (left->_left)
      {
        left = left->_left;
      }
 
      _node = left;
    }
    else // 右子树不存在，++后就到：所在子树是左子树的最近祖宗节点。
    {
      Node* parent = _node->_parent; // 找祖先里面孩子不是祖先的右的那个
      Node* cur = _node;
      while (parent && cur == parent->_right)
      {   // 是右子树就往上走，parent存在是处理_node是中序最后一个结点的情况
        cur = cur->_parent;
        parent = parent->_parent;
      }
 
      _node = parent;
    }
 
    return *this;
  }
 
  Self operator++(int)
  {
    Self ret = Self(_node); // 记录当前位置 最后返回
 
    if (_node->_right) // 右子树存在，++后就到右子树的最左结点
    {
      Node* left = _node->_right;
      while (left->_left)
      {
        left = left->_left;
      }
 
      _node = left;
    }
    else // 右子树不存在，++后就到：所在子树是左子树的最近祖宗节点。
    {
      Node* parent = _node->_parent; // 找祖先里面孩子不是祖先的右的那个
      Node* cur = _node;
      while (parent && cur == parent->_right)
      {   // 是右子树就往上走，parent存在是处理_node是中序最后一个结点的情况
        cur = cur->_parent;
        parent = parent->_parent;
      }
 
      _node = parent;
    }
 
    return ret;
  }
 
  Self& operator--()
  {
    if (_node->_left) // 左子树存在，++后就到左子树的最右结点
    {
      Node* right = _node->_left;
      while (right->_right)
      {
        right = right->_right;
      }
 
      _node = right;
    }
    else // 左子树不存在，++后就到：所在子树是右子树的最近祖宗节点。
    {
      Node* parent = _node->_parent; // 找祖先里面孩子不是祖先的左的那个
      Node* cur = _node;
      while (parent && cur == parent->_left)
      {  // 是左子树就往上走，parent存在是处理_node是中序第一个结点的情况
        cur = cur->_parent;
        parent = parent->_parent;
      }
 
      _node = parent;
    }
 
    return *this;
  }
 
  Self operator--(int)
  {
    Self ret = Self(_node); // 记录当前位置 最后返回
 
    if (_node->_left) // 左子树存在，++后就到左子树的最右结点
    {
      Node* right = _node->_left;
      while (right->_right)
      {
        right = right->_right;
      }
 
      _node = right;
    }
    else // 左子树不存在，++后就到：所在子树是右子树的最近祖宗节点。
    {
      Node* parent = _node->_parent; // 找祖先里面孩子不是祖先的左的那个
      Node* cur = _node;
      while (parent && cur == parent->_left)
      {  // 是左子树就往上走，parent存在是处理_node是中序第一个结点的情况
        cur = cur->_parent;
        parent = parent->_parent;
      }
 
      _node = parent;
    }
 
    return ret;
  }
};
 
template<class K, class T, class KeyOfT> // KeyOfT仿函数，把key取出来
struct RBTree
{
  typedef RBTreeNode<T> Node;
public:
  typedef __RBTreeIterator<T, T&, T*> iterator;
 
  iterator begin()
  {
    Node* left = _root;
    while (left && left->_left)
    {
      left = left->_left;
    }
 
    return iterator(left);
  }
 
  iterator end()
  {
    return iterator(nullptr);
  }
 
  pair<iterator, bool> Insert(const T& data) // 1. 返回值改成pair<iterator, bool>
  {
    KeyOfT kot; // 定义一个仿函数的对象
 
    if (_root == nullptr)
    {
      _root = new Node(data);
      _root->_col = BLACK; // 根给黑色
      return make_pair(iterator(_root), true); // 2. 空树时返回根的迭代器和true的键值对
    }
 
    Node* parent = nullptr;
    Node* cur = _root;
    while (cur) // 找到要插入的结点
    {
      if (kot(cur->_data) < kot(data))
      {
        parent = cur;
        cur = cur->_right;
      }
      else if (kot(cur->_data) > kot(data))
      {
        parent = cur;
        cur = cur->_left;
      }
      else
      {
        return make_pair(cur, false); // 3. 存在新插入的数据，返回原本存在的数据的迭代器和false
      }
    }
 
    cur = new Node(data);
    Node* newnode = cur; // 记录最后插入成功返回新插入结点的迭代器
    cur->_col = RED; // 默认插入红色结点
    if (kot(parent->_data) < kot(data)) // 找到位置后插入结点
    {
      parent->_right = cur;
    }
    else
    {
      parent->_left = cur;
    }
    cur->_parent = parent;
 
    while (parent && parent->_col == RED) // 父亲存在且为红才需要处理
    {
      Node* grandfather = parent->_parent;
      assert(grandfather); // 确定的可以断言下，否则就是插入前就不是红黑树
      assert(grandfather->_col == BLACK);
      if (grandfather->_left == parent)
      {
        Node* uncle = grandfather->_right;
 
        if (uncle && uncle->_col == RED) // 情况一，叔叔存在且为红(可以直接复制到下面uncle在左边)
        {    // 将父亲和叔叔改为黑，祖父改为红，然后把祖父当成cur，parent变祖父parent继续向上调整。
          parent->_col = uncle->_col = BLACK;
          grandfather->_col = RED;
          cur = grandfather;
          parent = cur->_parent;
        }
        else // 情况二或情况三：叔叔存在且为黑或叔叔不存在
        {
          if (cur == parent->_left) // 情况二的右旋+变色(parent在左)
          {
            //     g      
            //   p   u
            // c
            RotateR(grandfather);
            parent->_col = BLACK; // 父亲变为根了
            grandfather->_col = RED;
          }
          else // 情况二的左右双旋+变色(parent在左)
          {
            //      g      
            //   p     u
            //    c
            RotateL(parent);
            RotateR(grandfather);
            cur->_col = BLACK; // cur变为根了
            grandfather->_col = RED;
          }
          break;
        }
      }
      else
      {
        Node* uncle = grandfather->_left;
 
        if (uncle && uncle->_col == RED) // 情况一，叔叔存在且为红
        {    // 将父亲和叔叔改为黑，祖父改为红，然后把祖父当成cur，parent变祖父parent继续向上调整。
          parent->_col = uncle->_col = BLACK;
          grandfather->_col = RED;
          cur = grandfather;
          parent = cur->_parent;
        }
        else // 情况二或情况三：叔叔存在且为黑或叔叔不存在
        {
          if (cur == parent->_right) // 情况二的左旋+变色(parent在右)
          {
            //     g      
            //   u   p
            //        c
            RotateL(grandfather);
            parent->_col = BLACK; // 父亲变为根了
            grandfather->_col = RED;
          }
          else // 情况二的右左双旋+变色(parent在右)
          {
            //       g      
            //    u     p
            //         c
            RotateR(parent);
            RotateL(grandfather);
            cur->_col = BLACK; // cur变为根了
            grandfather->_col = RED;
          }
          break;
        }
      }
    }
 
    _root->_col = BLACK;
    return make_pair(newnode, true); // 4. 插入成功，返回新插入数据的迭代器和true
  }
 
  void InOrder()
  {
    _InOrder(_root);
    cout << endl;
  }
 
  bool IsBalance()
  {
    if (_root == nullptr)
    {
      return true;
    }
 
    if (_root->_col == RED) // 验证性质二
    {
      cout << "根节点不是黑色" << endl;
      return false;
    }
 
    int benchmark = 0; // 黑色节点数量基准值
    //Node* cur = _root; // 这种方法是先遍历一遍，然后传值，不过我们可以传引用
    //while (cur)
    //{
    //  if (cur->_col == BLACK)
    //  {
    //    ++benchmark;
    //  }
    //  cur = cur->_left;
    //}
    return PrevCheck(_root, 0, benchmark); // 验证性质三和四
  }
 
protected:
  bool PrevCheck(Node* root, int blackNum, int& benchmark)
  {
    if (root == nullptr)
    {
      if (benchmark == 0)
      {
        benchmark = blackNum;
        return true;
      }
 
      if (blackNum != benchmark) // 验证性质三
      {
        cout << "某条黑色节点的数量不相等" << endl;
        return false;
      }
      else
      {
        return true;
      }
    }
 
    if (root->_col == BLACK)
    {
      ++blackNum;
    }
 
    if (root->_col == RED && root->_parent->_col == RED) // 验证性质四
    {
      cout << "存在连续的红色节点" << endl;
      return false;
    }
 
    return PrevCheck(root->_left, blackNum, benchmark)
      && PrevCheck(root->_right, blackNum, benchmark);
  }
 
  void _InOrder(Node* root)
  {
    if (root == nullptr)
    {
      return;
    }
    _InOrder(root->_left);
    cout << root->_kv.first << ":" << root->_kv.second << endl;
    _InOrder(root->_right);
  }
 
  void RotateL(Node* parent)
  {
    Node* subR = parent->_right; // 动了三个标记了的结点，共更新六个指针，这更新两个指针
    Node* subRL = subR->_left;
 
    parent->_right = subRL;
    if (subRL) // subRL不为空才更新
    {
      subRL->_parent = parent;
    }
 
    Node* ppNode = parent->_parent; // 记录parent的parent，防止parent是一颗子树的头结点
 
    subR->_left = parent; // 再更新两个指针
    parent->_parent = subR;
 
    if (_root == parent)  // 最后更新两个指针
    {
      _root = subR;
      subR->_parent = nullptr;
    }
    else // parent是一颗子树的头结点
    {
      if (ppNode->_left == parent)
      {
        ppNode->_left = subR;
      }
      else
      {
        ppNode->_right = subR;
      }
      subR->_parent = ppNode;
    }
  }
 
  void RotateR(Node* parent)
  {
    Node* subL = parent->_left;
    Node* subLR = subL->_right;
 
    parent->_left = subLR; // 更新两个节点
    if (subLR)
    {
      subLR->_parent = parent;
    }
 
    Node* ppNode = parent->_parent;
 
    subL->_right = parent; // 再更新两个节点
    parent->_parent = subL;
 
    if (_root == parent) // 最后更新两个结点
    {
      _root = subL;
      subL->_parent = nullptr;
    }
    else
    {
      if (ppNode->_left == parent)
      {
        ppNode->_left = subL;
      }
      else
      {
        ppNode->_right = subL;
      }
 
      subL->_parent = ppNode;
    }
  }
 
  Node* _root = nullptr;
};