深入红黑树：从STL源码到亲手封装 map 和 set

FeIix

404人浏览 · 2026-05-10 21:50:35

FeIix · 2026-05-10 21:50:35 发布

引言

在学习 C++ 标准模板库时，map 和 set 是我们最常用的关联式容器。你是否想过，为什么 map 存储的是键值对，而 set 只存储键——但它们却共用同一套底层数据结构？为什么它们的迭代器都是有序的？

1. 源码框架层：一棵树如何撑起两种容器

1.1 核心矛盾

打开 SGI STL 源码，stl_map.h 和 stl_set.h 中都包含了 stl_tree.h。它们底层都依赖同一棵红黑树，但问题来了：

set 存放的是单纯的 Key
map 存放的是 pair<const Key, T>

如何让同一棵树既存 Key 又存 pair？

1.2 泛型模板参数的艺术

观察源码中 rb_tree 的模板参数：

template <class Key, class Value, class KeyOfValue, class Compare, class Alloc = alloc>
class rb_tree { ... };

关键在第二个模板参数 Value：

set 实例化时：rb_tree<key_type, key_type, identity<value_type>, ...>
- 第二个参数直接传 key_type，树节点存储的就是 Key 本身
map 实例化时：rb_tree<key_type, pair<const Key, T>, select1st<value_type>, ...>
- 第二个参数传 pair<const Key, T>，树节点存储键值对

这样，一棵红黑树通过模板参数 Value 控制了节点存储的真实数据类型，完美复用于两种场景。

1.3 三个模板参数的分工

模板参数	set 中	map 中	作用
Key	Key	Key	`erase` / `find` 的形参类型
Value	Key	`pair<const Key, T>`	决定节点存什么数据
KeyOfValue	`identity<Key>`	`select1st<pair>`	从节点数据中提取 Key

为什么需要 KeyOfValue？
因为插入时需要比较大小，但 pair 默认的比较会先比 key 再比 value——而我们只希望按 key 比较。所以需要仿函数从 Value 中“取出”Key 来比较。

2. 动手实现： mymap 和 myset

2.1 调整模板命名

相比源码，我们做一点命名优化：

第一个参数用 K（对应源码的 Key）
第二个参数用 V（源码中 map 的第二个参数叫 T，容易混淆）
红黑树中的数据类型用 T

2.2 实现 KeyOfT 仿函数

Mymap.h：

template<class K, class V>
class map
{
    struct MapKeyOfT
    {
        const K& operator()(const pair<K, V>& kv)
        {
            return kv.first;  // 从 pair 中提取 key
        }
    };
public:
    bool insert(const pair<K, V>& kv)
    {
        return _t.Insert(kv);
    }
private:
    RBTree<K, pair<K, V>, MapKeyOfT> _t;
};

Myset.h：

template<class K>
class set
{
    struct SetKeyOfT
    {
        const K& operator()(const K& key)
        {
            return key;  // set 的 key 就是自己
        }
    };
public:
    bool insert(const K& key)
    {
        return _t.Insert(key);
    }
private:
    RBTree<K, K, SetKeyOfT> _t;
};

这样，红黑树内部通过 KeyOfT 仿函数统一提取 Key 进行比较，完全不关心 T 到底是什么类型。

3. 迭代器实现：中序遍历的巧妙抽象

3.1 大框架：封装节点指针

迭代器的实现思路与 list 的迭代器一致——封装一个节点指针，重载 *、->、++、-- 等运算符，使其像指针一样访问数据。

真正的难点是 operator++ 和 operator--。

3.2 operator++ 的逻辑

红黑树的迭代器按中序遍历（左子树 → 根 → 右子树）移动：

情况一：右子树不为空
当前节点访问完毕，下一个节点是右子树的最左节点。

情况二：右子树为空
说明当前子树全部访问完毕，需要沿着祖先路径向上找：

如果当前节点是父节点的左孩子，则下一个就是父节点
如果当前节点是父节点的右孩子，则继续往上找，直到找到孩子是父亲左孩子的那个祖先

Self& operator++()
{
    if (_node->_right)
    {
        // 右不为空，找右子树的最左节点
        Node* leftMost = _node->_right;
        while (leftMost->_left)
            leftMost = leftMost->_left;
        _node = leftMost;
    }
    else
    {
        // 右为空，找孩子是父亲左的那个祖先
        Node* cur = _node;
        Node* parent = cur->_parent;
        while (parent && cur == parent->_right)
        {
            cur = parent;
            parent = cur->_parent;
        }
        _node = parent;
    }
    return *this;
}

3.3 end() 的设计

我们用 nullptr 充当 end()。当迭代器指向最右节点（最大节点）再 ++ 时，找不到”孩子是父亲左”的祖先，最终 parent 为空，_node 被置为 nullptr，即 end()。

--end() 需要特殊处理：

Self& operator--()
{
    if (_node == nullptr) // end()
    {
        // 回到中序最后一个节点（整棵树的最右节点）
        Node* rightMost = _root;
        while (rightMost && rightMost->_right)
            rightMost = rightMost->_right;
        _node = rightMost;
    }
    // ... 后续逻辑与 ++ 对称
}

4. 进阶细节：const 支持和 operator[]

4.1 key 不可修改的问题

set：值就是 key，所以整个值不能修改。把 set 的第二个模板参数改为 const K：

RBTree<K, const K, SetKeyOfT> _t;

map：key 不能修改但 value 可以。把 pair 的第一个参数改为 const K：

RBTree<K, pair<const K, V>, MapKeyOfT> _t;

4.2 map 的 operator[]

改进的Insert

template<class Key, class T, class KeyOfT>
pair<typename RBTree<Key, T, KeyOfT>::Iterator,bool> RBTree<Key, T, KeyOfT>::Insert(const T& data)
{
	if (_root == nullptr)
	{
		_root = new Node(data);
		_root->_col = Black;
		return pair<Iterator, bool>({ Iterator(_root, _root),true });
	}

	Node* parent = nullptr;
	Node* cur = _root;

	KeyOfT koft;
	while (cur)
	{
		if (koft(cur->_data) < koft(data))
		{
			parent = cur;
			cur = cur->_right;
		}
		else if(koft(cur->_data) > koft(data))
		{
			parent = cur;
			cur = cur->_left;
		}else
		{
			return pair<Iterator, bool>({ Iterator(cur,_root),false});
		}
	}

	cur = new Node(data);
	if (koft(data) < koft(parent->_data))
	{
		parent->_left = cur;
	}
	else
	{
		parent->_right = cur;
	}
	cur->_parent = parent;

	while (parent && parent->_col == Red)
	{
		Node* grandfather = parent->_parent;
		if (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;
				}
				else
				{
					_rotateL(parent);
					_rotateR(grandfather);
					cur->_col = Black;
				}
				grandfather->_col = Red;
				break;
			}
		}
		else
		{
			Node* uncle = parent->_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;
				}
				else
				{
					_rotateR(parent);
					_rotateL(grandfather);
					cur->_col = Black;
				}
				grandfather->_col = Red;
				break;
			}
			
		}
	}

	_root->_col = Black;
	return pair<Iterator, bool>({ Iterator(cur,_root),true });
}

有了改进后的 Insert（返回 pair<iterator, bool>），operator[] 实现极简：

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

如果 key 已存在，insert 返回已有节点的迭代器；不存在则插入 key: V() 并返回新节点迭代器。最终都返回对应 value 的引用。

5. 总结：泛型思想的胜利

通过封装红黑树实现 map 和 set，我们看到泛型编程的威力：

层次	职责
RBTree	数据结构核心：旋转、调整平衡、迭代器遍历
KeyOfT	策略层：从不同类型数据中提取 Key 用于比较
mymap / myset	适配层：向用户暴露接口，指定存储类型和只读策略