【C++高阶】高效数据存储：理解并模拟实现红黑树Map与Set

// Map 与 Set
// set -> RBTree<K, K, SetKeyOfT> _t;
// map -> RBTree<K, pair<const K, V, MapKeyOfT>> _t;

enum Color
{
	RED,
	BLACK
};

template<class T>
struct RBTreeNode
{
	RBTreeNode<T>* _left;
	RBTreeNode<T>* _right;
	RBTreeNode<T>* _parent;
	// pair<K, V> _kv; // 上节内容使用的这种方式
	T _data;
	Color _col;

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

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

		Node* parent = nullptr;
		Node* cur = _root;
		// 通过仿函数KeyOfT来比较数据的大小
		KeyOfT kot;

		while (cur)
		{
			parent = cur;
			if (kot(cur->_data) < kot(data))
			{
				cur = cur->_right;
			}
			else if (kot(cur->_data) > kot(data))
			{
				cur = cur->_left;
			}
			else
			{
				return make_pair(cur, false);
			}
		}

		// 新增节点给红色
		cur = new Node(data);
		// 注意，这里需要保存一个新增节点，以便用来返回
		Node* newnode = cur;
		cur->_col = RED;
		// 旋转变色
		return make_pair(newnode, true);;
	}

// 通过模板来达到const的迭代器的复用
template<class T, class Ref, class Ptr>
struct __TreeIterator
{
	typedef RBTreeNode<T> Node;
	typedef __TreeIterator<T, Ref, Ptr> Self;
	Node* _node;
	
	// 构造函数
	__TreeIterator(Node* node)
		:_node(node)
	{}

	Ref operator*()
	{
		return _node->_data;
	}

	Ptr operator->()
	{
		return &(_node->_data);
	}

	bool operator!=(const Self& s)
	{
		return _node != s._node;
	}
};

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* cur = _node->_right;
		while (cur->_left)
		{
			cur = cur->_left;
		}
		_node = cur;
	}
	else
	{
		Node* cur = _node;
		Node* parent = cur->_parent;
		while (parent && cur == parent->_right)
		{
			cur = parent;
			parent = cur->_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 = cur->_parent;
		}
		_node = parent;
	}
	return *this;
}

template<class K>
class set
{
public:
	// SetKeyOfT：通过T来获取key的一个仿函数类
	struct SetKeyOfT
	{
		const K& operator()(const K& key)
		{
			return key;
		}
	};
	// 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();
	}
	
	// 复用RBTree中的插入
	pair<iterator, bool> insert(const K& key)
	{
		return _t.Insert(key);
	}

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

template<class K, class V>
class map
{
public:
	// MapKeyOfT：通过pair来获取kv.first的一个仿函数类
	struct MapKeyOfT
	{
		const K& operator()(const pair<K, V>& kv)
		{
			return kv.first;
		}
	};
	
	// map是一个key不能修改，value能够修改的结构
	typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::iterator iterator;
	typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::iterator const_iterator;

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

	iterator end()
	{
		return _t.end();
	}
	
	// 重载map中的operatorp[]
	// operatorp[] 当原数据中没有时，插入并初始化，有则修改second
	V& operator[](const K& key)
	{
		// 没有是插入，有则是查找
		pair<iterator, bool> ret = insert(make_pair(key, V()));
		return ret.first->second;
	}
	
	// 复用RBTree中的插入
	pair<iterator, bool> insert(const pair<K, V>& kv)
	{
		return _t.Insert(kv);
	}
	
private:
	RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};