# 从红黑树到标准容器：手把手实现STL Set/Map核心机制
红黑树作为平衡二叉搜索树的经典实现，是C++标准库中set和map容器的底层数据结构。理解红黑树的实现原理不仅有助于深入掌握STL容器的内部机制，还能提升算法设计和内存管理能力。本文将带领你从零开始实现一棵红黑树，并逐步封装成符合STL风格的set容器。
## 红黑树基础：理解五大约束条件
红黑树是一种自平衡二叉搜索树，它通过以下五条规则维持平衡：
1. 每个节点要么是红色，要么是黑色
2. 根节点是黑色
3. 所有叶子节点（NIL节点）都是黑色
4. 红色节点的两个子节点都是黑色
5. 从任一节点到其每个叶子节点的所有路径包含相同数目的黑色节点
```cpp
enum Color { RED, BLACK };
template<typename T>
struct RBTreeNode {
T data;
Color color;
RBTreeNode* left;
RBTreeNode* right;
RBTreeNode* parent;
RBTreeNode(const T& val, Color c = RED)
: data(val), color(c), left(nullptr),
right(nullptr), parent(nullptr) {}
};
```
## 红黑树核心操作实现
### 旋转操作：维持平衡的基础
旋转操作是红黑树维持平衡的关键，包括左旋和右旋两种基本操作：
```cpp
template<typename T>
class RedBlackTree {
private:
RBTreeNode<T>* root;
RBTreeNode<T>* nil; // 哨兵节点
// 左旋操作
void leftRotate(RBTreeNode<T>* x) {
RBTreeNode<T>* y = x->right;
x->right = y->left;
if (y->left != nil) {
y->left->parent = x;
}
y->parent = x->parent;
if (x->parent == nil) {
root = y;
} else if (x == x->parent->left) {
x->parent->left = y;
} else {
x->parent->right = y;
}
y->left = x;
x->parent = y;
}
// 右旋操作（与左旋对称）
void rightRotate(RBTreeNode<T>* y) {
RBTreeNode<T>* x = y->left;
y->left = x->right;
if (x->right != nil) {
x->right->parent = y;
}
x->parent = y->parent;
if (y->parent == nil) {
root = x;
} else if (y == y->parent->left) {
y->parent->left = x;
} else {
y->parent->right = x;
}
x->right = y;
y->parent = x;
}
};
```
### 插入操作的平衡调整
插入新节点后，可能破坏红黑树的性质，需要通过调整颜色和旋转来恢复平衡：
```cpp
template<typename T>
void RedBlackTree<T>::insertFixup(RBTreeNode<T>* z) {
while (z->parent->color == RED) {
if (z->parent == z->parent->parent->left) {
RBTreeNode<T>* y = z->parent->parent->right;
// 情况1：叔节点为红色
if (y->color == RED) {
z->parent->color = BLACK;
y->color = BLACK;
z->parent->parent->color = RED;
z = z->parent->parent;
} else {
// 情况2：z是右孩子
if (z == z->parent->right) {
z = z->parent;
leftRotate(z);
}
// 情况3：z是左孩子
z->parent->color = BLACK;
z->parent->parent->color = RED;
rightRotate(z->parent->parent);
}
} else {
// 对称情况处理
RBTreeNode<T>* y = z->parent->parent->left;
if (y->color == RED) {
z->parent->color = BLACK;
y->color = BLACK;
z->parent->parent->color = RED;
z = z->parent->parent;
} else {
if (z == z->parent->left) {
z = z->parent;
rightRotate(z);
}
z->parent->color = BLACK;
z->parent->parent->color = RED;
leftRotate(z->parent->parent);
}
}
}
root->color = BLACK;
}
```
## 迭代器设计：符合STL规范
实现STL风格的容器必须提供符合要求的迭代器：
```cpp
template<typename T>
class RBIterator {
private:
RBTreeNode<T>* current;
RBTreeNode<T>* nil;
public:
using iterator_category = std::bidirectional_iterator_tag;
using value_type = T;
using difference_type = std::ptrdiff_t;
using pointer = T*;
using reference = T&;
RBIterator(RBTreeNode<T>* node, RBTreeNode<T>* sentinel)
: current(node), nil(sentinel) {}
// 解引用操作符
reference operator*() const {
return current->data;
}
pointer operator->() const {
return &(current->data);
}
// 前向迭代
RBIterator& operator++() {
if (current->right != nil) {
current = minimum(current->right);
} else {
RBTreeNode<T>* parent = current->parent;
while (parent != nil && current == parent->right) {
current = parent;
parent = parent->parent;
}
current = parent;
}
return *this;
}
// 反向迭代
RBIterator& operator--() {
if (current == nil) {<"kw.dfkngj.com">
current = maximum(root);
} else if (current->left != nil) {
current = maximum(current->left);
} else {
RBTreeNode<T>* parent = current->parent;
while (parent != nil && current == parent->left) {
current = parent;
parent = parent->parent;
}
current = parent;
}
return *this;
}
private:
RBTreeNode<T>* minimum(RBTreeNode<T>* node) {
while (node->left != nil) {
node = node->left;
}
return node;
}
RBTreeNode<T>* maximum(RBTreeNode<T>* node) {
while (node->right != nil) {
node = node->right;
}
return node;
}
};
```
## 封装为STL风格的Set容器
基于红黑树实现符合STL接口规范的set容器：
```cpp
template<typename Key,
typename Compare = std::less<Key>,
typename Allocator = std::allocator<Key>>
class my_set {
private:
RedBlackTree<Key> rbtree;
Compare comp;
Allocator alloc;
size_t element_count;
public:
// 类型定义
using key_type = Key;
using value_type = Key;
using size_type = size_t;
using difference_type = std::ptrdiff_t;
using key_compare = Compare;
using value_compare = Compare;
using allocator_type = Allocator;
using reference = value_type&;
using const_reference = const value_type&;
using pointer = typename std::allocator_traits<Allocator>::pointer;
using const_pointer = typename std::allocator_traits<Allocator>::const_pointer;
using iterator = RBIterator<value_type>;
using const_iterator = RBIterator<const value_type>;
// 构造函数
my_set() : element_count(0) {}
explicit my_set(const Compare& comp,
const Allocator& alloc = Allocator())
: comp(comp), alloc(alloc), element_count(0) {}
template<typename InputIt><"rz.dfkngj.com">
my_set(InputIt first, InputIt last,
const Compare& comp = Compare(),
const Allocator& alloc = Allocator())
: comp(comp), alloc(alloc), element_count(0) {
insert(first, last);
}
// 容量操作
bool empty() const noexcept { return element_count == 0; }
size_type size() const noexcept { return element_count; }
// 修改操作
std::pair<iterator, bool> insert(const value_type& value) {
bool inserted = rbtree.insert(value);
if (inserted) {
++element_count;
}
return {find(value), inserted};
}
iterator erase(iterator pos) {
iterator next = pos;
++next;
if (rbtree.erase(*pos)) {
--element_count;
}
return next;
}
// 查找操作
iterator find(const key_type& key) {
return iterator(rbtree.search(key), rbtree.getNil());
}
const_iterator find(const key_type& key) const {
return const_iterator(rbtree.search(key), rbtree.getNil());
}
// 迭代器支持
iterator begin() noexcept {
return iterator(rbtree.minimum(), rbtree.getNil());
}
iterator end() noexcept {
return iterator(rbtree.getNil(), rbtree.getNil());
}
const_iterator begin() const noexcept {
return const_iterator(rbtree.minimum(), rbtree.getNil());
}
const_iterator end() const noexcept {
return const_iterator(rbtree.getNil(), rbtree.getNil());
}
// 边界查找
iterator lower_bound(const key_type& key) {
return iterator(rbtree.lower_bound(key), rbtree.getNil());
}
iterator upper_bound(const key_type& key) {
return iterator(rbtree.upper_bound(key), rbtree.getNil());
}
};
```
## 扩展为Map容器
基于相同的红黑树结构，我们可以实现map容器，主要区别在于存储键值对：
```cpp
template<typename Key, typename T><"mv.dfkngj.com">
struct Pair {
Key first;
T second;
bool operator<(const Pair& other) const {
return first < other.first;
}
bool operator==(const Pair& other) const {
return first == other.first;
}
};
template<typename Key, typename T,
typename Compare = std::less<Key>,
typename Allocator = std::allocator<Pair<Key, T>>>
class my_map {
private:
RedBlackTree<Pair<Key, T>> rbtree;
size_t element_count;
public:
using key_type = Key;
using mapped_type = T;
using value_type = Pair<Key, T>;
// 重载[]操作符
T& operator[](const key_type& key) {
auto it = find(key);
if (it == end()) {
it = insert({key, T()}).first;
}
return it->second;
}
// 插入键值对
std::pair<iterator, bool> insert(const value_type& value) {
return rbtree.insert(value);
}
// 范围插入
template<typename InputIt>
void insert(InputIt first, InputIt last) {
while (first != last) {
insert(*first);
++first;
}
}
};
```
## 性能分析与优化
### 时间复杂度保证
通过红黑树的平衡特性，我们的实现保证了以下时间复杂度：
- 插入：O(log n)
- 删除：O(log n)
- 查找：O(log n)
- 遍历：O(n)
### 内存管理优化
```cpp
// 自定义内存分配器示例
template<typename T>
class PoolAllocator {
private:
std::vector<char*> memory_pools;
size_t pool_size;
size_t object_size;
public:
pointer allocate(size_type n) {
if (n != 1) {
return static_cast<pointer>(::operator new(n * sizeof(T)));
}
// 实现对象池分配逻辑
// ...
}
void deallocate(pointer p, size_type n) {
if (n != 1) {
::operator delete(p);
return;
}
// 实现对象池回收逻辑
// ...
}
};
```
## 测试验证
完整的容器实现需要严格的测试验证：
```cpp
void test_my_set() {
my_set<int> s;
// 插入测试
for (int i = 0; i < 1000; ++i) {
s.insert(i * 2);
}
// 查找测试
assert(s.find(500) != s.end());
assert(s.find(501) == s.end());
// 遍历测试
int count = 0;
for (auto it = s.begin(); it != s.end(); ++it) {
++count;
}
assert(count == 1000);
// 删除测试
s.erase(s.find(100));
assert(s.find(100) == s.end());
assert(s.size() == 999);
}
```
## 总结与扩展
通过从红黑树基础实现到完整STL容器的封装，我们深入理解了平衡二叉搜索树的工作原理和标准容器的设计哲学。这种实现方式不仅有助于理解STL的内部机制，也为自定义数据结构提供了参考模式。
可以进一步扩展的功能包括：
1. 支持自定义分配器
2. 实现异常安全保证
3. 添加移动语义支持
4. 实现线程安全版本
5. 支持初始化列表构造
理解这些底层实现原理，能够帮助开发者在实际工作中更好地使用标准容器，并在需要时进行适当的扩展和优化。
