在代码的红与黑间——红黑树实现 map 和 set 的美丽旅程

enum Color {
    RED,
    BLACK
};

template <class T>
class RBTreeNode {
public:
    RBTreeNode<T>* left;
    RBTreeNode<T>* right;
    RBTreeNode<T>* _parent;
    T _data;
    Color _col;

    RBTreeNode(const T& data) 
        : left(nullptr), right(nullptr), _parent(nullptr), _data(data), _col(RED) {}
};

pair<Iterator, bool> insert(const T& data) {
	if (!_root) {
		_root = new Node(data);
		_root->_col = BLACK;
		return make_pair(Iterator(_root), true);
	}

	// 定义当前节点cur
	Node* cur = _root;
	Node* parent = nullptr;

	// 仿函数用来索引key
	KeyOfT kot;
	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(Iterator(cur), false); // 不可能出现相等的情况插入
	}

	// 创建新的Node实例cur，并将它插入到适当的位置
	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;
		// 父亲是祖父的左孩子
		if (grandfather && parent == grandfather->left) {
			Node* uncle = grandfather->right;
			// 叔叔存在且为红
			if (uncle && uncle->_col == RED) {
				parent->_col = uncle->_col = BLACK;
				grandfather->_col = RED;

				//继续向上处理
				cur = grandfather;
				parent = cur->_parent;
			}
			// 叔叔不存在或者存在且为黑
			else {
				if (cur == parent->left) {
					RotateR(grandfather);
					parent->_col = BLACK;
					grandfather->_col = RED;
				}
				else {
					RotateL(parent);
					RotateR(grandfather);
					cur->_col = BLACK;
					grandfather->_col = RED;
				}

				break;
			}
		}
		// 父亲是祖父的右孩子
		else if (grandfather && parent == grandfather->right) {
			Node* uncle = grandfather->left;
			// 叔叔存在且为红
			if (uncle && uncle->_col == RED) {
				parent->_col = uncle->_col = BLACK;
				grandfather->_col = RED;

				//继续向上处理
				cur = grandfather;
				parent = cur->_parent;
			}
			// 叔叔不存在或者存在且为黑
			else {
				if (cur == parent->right) {
					RotateL(grandfather);
					parent->_col = BLACK;
					grandfather->_col = RED;
				}
				else {
					RotateR(parent);
					RotateL(grandfather);
					cur->_col = BLACK;
					grandfather->_col = RED;
				}

				break;
			}
		}
	}
	_root->_col = BLACK;

	return make_pair(Iterator(newnode), true);
}

template<class T, class Ref, class Ptr>
class __TreeIterator {
public:
    typedef RBTreeNode<T> Node;
    typedef __TreeIterator<T, Ref, Ptr> Self;
    typedef __TreeIterator<T, T&, T*> Iterator;
    Node* node;
	
    // 如果迭代器是const，则为构造函数，否则为拷贝构造函数
    __TreeIterator(const Iterator& it) : node(it._node) {}
    __TreeIterator(Node* node = nullptr) : node(node) {}

    Ref operator*() { return node->data; }
    Ptr operator->() { return &(node->data); }

    bool operator!=(const Self& other) const { return node != other.node; }

    Self& operator++();
    Self& operator--();
};

typedef __TreeIterator<T, T&, T*> Iterator;
typedef __TreeIterator<T, const T&, const T*> const_Iterator;

Iterator begin() {
    Node* cur = root;
    while (cur && cur->left) {
        cur = cur->left;
    }
    return Iterator(cur);
}

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

const_Iterator begin() const {
	Node* cur = _root;
	while (cur && cur->left) {
		cur = cur->left;
	}
	return  const_Iterator(cur);
}

const_Iterator end() const {
	return const_Iterator(nullptr);
}

Self& operator++() {
    if (node->right) {
        node = node->right;
        while (node->left) {
            node = node->left;
        }
    } else {
        Node* parent = node->parent;
        while (parent && node == parent->right) {
            node = parent;
            parent = parent->parent;
        }
        node = parent;
    }
    return *this;
}

Self& operator--() {
	if (_node->left) {
		Node* cur = _node->left;
		while (cur->right) {
			cur = cur->right;
		}
		_node = cur;
	}
	else {
		Node* cur = _node;
		Node* parent = cur->_parent;
		while (parent && cur == parent->left) {
			cur = parent;
			parent = parent->_parent;
		}
		_node = parent;
	}

	return *this;
}

class SetKeyOfT {
public:
    const K& operator()(const K& key){
        return key;
    }
};

template<class K>
class set{
public:
    // typename告诉编译器这是类型，
	// 因为set不能够进行修改，所以我们用const迭代器来初始化普通迭代器，来达到不能修改的目的
    typedef typename RBTree<K, K, SetKeyOfT>::const_Iterator Iterator;
    typedef typename RBTree<K, K, SetKeyOfT>::const_Iterator const_Iterator;

    Iterator begin() const {
        return _t.begin();
    }

    Iterator end() const {
        return _t.end();
    }

    pair<Iterator, bool> insert(const K& key) {
        pair<typename RBTree<K, K, SetKeyOfT>::Iterator, bool> ret = _t.insert(key);
        return pair<Iterator, bool>(ret.first, ret.second);
    }

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

class MapKeyOfT {
public:
    const K& operator()(const pair<K, V>& kv) {
        return kv.first;
    }
};

template<class K,class V>
class map {
public:
    // map是一个key不能修改，value能够修改的结构
    typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::Iterator Iterator;
    typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::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();
    }

    pair<Iterator, bool> insert(const pair<K, V>& kv) {
        return _t.insert(kv);
    }

    // 重载map中的operatorp[]
	// operatorp[] 当原数据中没有时，插入并初始化，有则修改second
    V& operator[](const K& key) {
        pair<Iterator, bool> ret = insert(make_pair(key, V()));
        return ret.first->second;
    }

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

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