红黑树封装map/set核心技巧

红黑树封装 map 和 set 的核心逻辑

红黑树是一种自平衡二叉搜索树，适合作为 map 和 set 的底层实现。map 存储键值对，set 只存储键，但两者都可以基于同一棵红黑树实现。

红黑树节点设计

红黑树节点需要包含颜色标记、左右子节点指针和父节点指针。对于 map 和 set 的封装，节点数据设计有所差异：

// 基础红黑树节点
enum Color { RED, BLACK };

template <typename T>
struct RBTreeNode {
    T data;
    Color color;
    RBTreeNode* left;
    RBTreeNode* right;
    RBTreeNode* parent;
};

map 和 set 的节点数据差异

对于 set，节点直接存储键值：

// set 的节点数据就是 Key 本身
template <typename Key>
using SetNode = RBTreeNode<Key>;

对于 map，节点存储键值对：

// map 的节点数据是 pair<const Key, Value>
template <typename Key, typename Value>
using MapNode = RBTreeNode<std::pair<const Key, Value>>;

迭代器设计

红黑树的迭代器需要支持中序遍历（有序遍历）。迭代器核心是维护当前节点的指针：

template <typename T>
class RBTIterator {
    RBTreeNode<T>* current;
public:
    // 前置++
    RBTIterator& operator++() {
        if (current->right) {
            current = current->right;
            while (current->left) current = current->left;
        } else {
            RBTreeNode<T>* p = current->parent;
            while (p && current == p->right) {
                current = p;
                p = p->parent;
            }
            current = p;
        }
        return *this;
    }
    // 解引用
    T& operator*() { return current->data; }
};

map 和 set 的接口差异

set 的接口主要围绕键操作：

template <typename Key>
class set {
    RBTree<Key> tree;
public:
    iterator find(const Key& key);
    void insert(const Key& key);
};

map 的接口需要处理键值对：

template <typename Key, typename Value>
class map {
    RBTree<std::pair<const Key, Value>> tree;
public:
    Value& operator[](const Key& key) {
        auto it = tree.find(key);
        if (it == end()) 
            it = tree.insert({key, Value()});
        return it->second;
    }
};

红黑树的插入操作封装

红黑树的插入需要处理颜色调整和旋转。插入操作对 map 和 set 是通用的：

template <typename T>
void RBTree<T>::insert(const T& value) {
    // 标准BST插入
    RBTreeNode<T>* newNode = new RBTreeNode<T>{value, RED};
    // ... 插入逻辑
    
    // 插入后修复红黑树性质
    insertFixup(newNode);
}

模板参数的设计技巧

通过模板参数实现代码复用：

template <typename Key, typename Value, typename Compare = less<Key>>
class map {
    using NodeType = RBTreeNode<std::pair<const Key, Value>>;
    using TreeType = RBTree<NodeType, Compare>;
    TreeType tree;
    // ... 接口实现
};

关键点总结

红黑树的节点设计需要同时满足 map 和 set 的需求。迭代器实现要考虑中序遍历。map 的 operator[] 需要特殊处理插入逻辑。模板参数的设计能最大化代码复用。

通过这种方式，可以用同一棵红黑树高效地实现 map 和 set 两种容器，同时保持各自的特性。这种设计模式在STL实现中广泛使用。