从零开始的 STL 实现记录：Container-6.4

前言

红黑树性质：

节点是红色或黑色。

根是黑色。

所有叶子都是黑色（叶子是 NIL 节点）。

每个红色节点必须有两个黑色的子节点。

从任一节点到其每个叶子的所有简单路径都包含相同数目的黑色节点。

连一刻都没有为国庆假期悼念，立刻赶来战场的是 ———— 长达七天的调休。而我们也将继续 STL 的旅程，本节将要分析的是 STL 红黑树的构造以及对外接口的实现。

rb-tree 的构造与内存管理

红黑树的构造与析构函数如下：

public:  // 构造、复制、析构函数
    rb_tree() { rb_tree_init(); }

    rb_tree(const rb_tree& rhs);
    rb_tree(rb_tree&& rhs) noexcept;

    rb_tree& operator=(const rb_tree& rhs);
    rb_tree& operator=(rb_tree&& rhs) noexcept;

    ~rb_tree() { clear(); }

/// @brief 复制构造函数
template <class T, class Compare>
rb_tree<T, Compare>::rb_tree(const rb_tree& rhs) {
    rb_tree_init();
    if (rhs.node_count_ != 0) {
        root() = copy_from(rhs.root(), header_);
        leftmost() = rb_tree_min(root());
        rightmost() = rb_tree_max(root());
    }
    node_count_ = rhs.node_count_;
    key_comp_ = rhs.key_comp_;
}

/// @brief 移动构造函数
template <class T, class Compare>
rb_tree<T, Compare>::rb_tree(rb_tree&& rhs) noexcept 
    : header_(tinystl::move(rhs.header_)), 
    node_count_(rhs.node_count_), 
    key_comp_(rhs.key_comp_) {
    rhs.reset();
}

可以看出，除移动构造外，构造函数均调用了 rb_tree_init() 函数来初始化红黑树，让我们来看看这个函数的实现：

/// @brief 初始化 rb-tree
template <class T, class Compare>
void rb_tree<T, Compare>::rb_tree_init() {
    header_ = base_allocator::allocate(1);
    header_->color = rb_tree_red;  // header_ 为红色，与 root 区分
    root() = nullptr;
    leftmost() = header_;
    rightmost() = header_;
    node_count_ = 0;
}

在 Container-6.2 中也分析了 STL 红黑树中的特殊设计：header 节点，该节点的存在使得我们可以很轻松的处理各种边界情况：如判断是否查找完、找最大最小节点等。

因此，维护正确的 header 对于红黑树来说就至关重要。初始化函数会生成一个 header 节点并将其染为红色，以区分根节点与 header 节点，同时令 root() = nullptr、leftmost 与 rightmost 都指向自身。

由于 tinystl 中额外维护了 node_count 的信息，因此这里也要进行相应的初始化。

/// @brief 重置 rb-tree
template <class T, class Compare>
void rb_tree<T, Compare>::reset() {
    header_ = nullptr;
    node_count_ = 0;
}

在移动构造函数中用到了 reset() 这个函数，就是简单地将 header_ 与 node_count_ 归零，但值得注意的是，该函数只在移动构造与移动赋值中才能使用，否则会出现严重的内存泄漏问题。

/// @brief 清空 rb-tree
template <class T, class Compare>
void rb_tree<T, Compare>::clear() {
    if (node_count_ != 0) {
        erase_since(root());
        leftmost() = header_;
        root() = nullptr;
        rightmost() = header_;
        node_count_ = 0;
    }
}

另外，可以看出 rb-tree 的析构函数就是调用了 clear() 函数，分析该函数的实现可以发现，该函数先析构并释放所有的节点，之后再把 root、leftmost 、rightmost 以及 node_count_ 都重置一遍，关于 erase_since(root()) 的实现会在后面进行说明。

除了 rb_tree_init() 与 reset() 外，还有以下几个关于节点的构造与销毁的函数经常被使用到，这里一并列出。

/// @brief 创建节点
template <class T, class Compare>
template <class ...Args>
typename rb_tree<T, Compare>::node_ptr
rb_tree<T, Compare>::create_node(Args&& ...args) {
    auto tmp = node_allocator::allocate(1);
    try {
        // 在节点位置构造元素
        data_allocator::construct(tinystl::address_of(tmp->value), tinystl::forward<Args>(args)...);
        tmp->parent = nullptr;
        tmp->left = nullptr;
        tmp->right = nullptr;
    }
    catch (...) {
        node_allocator::deallocate(tmp);
        throw;
    }
    return tmp;
}

/// @brief 复制节点
template <class T, class Compare>
typename rb_tree<T, Compare>::node_ptr
rb_tree<T, Compare>::clone_node(base_ptr x) {
    // auto tmp = create_node(x->get_node_ptr()->value);
    auto tmp = create_node(static_cast<node_ptr>(x)->value);
    tmp->color = x->color;
    tmp->left = nullptr;
    tmp->right = nullptr;
    return tmp;
}

/// @brief 销毁节点
template <class T, class Compare>
void rb_tree<T, Compare>::destroy_node(node_ptr x) {
    data_allocator::destroy(tinystl::address_of(x->value));
    node_allocator::deallocate(x);
}

rb-tree 的元素插入

rb-tree 提供多种插入元素的操作，由于需要支持可重复的容器与不可重复的容器，因此需要两套不同的插入操作；同时在内部还需要处理插入元素为 pair 与普通元素各自的情况。因此较为繁琐。以下以 insert_multi() 代表可重复容器的插入操作，insert_unique() 代表不可重复容器的插入操作。

public:  // 元素相关操作

    // ====================== emplace ====================== //
    template <class ...Args>
    iterator  emplace_multi(Args&& ...args);

    template <class ...Args>
    tinystl::pair<iterator, bool> emplace_unique(Args&& ...args);

    template <class ...Args>
    iterator  emplace_multi_use_hint(iterator hint, Args&& ...args);

    template <class ...Args>
    iterator  emplace_unique_use_hint(iterator hint, Args&& ...args);

    // ====================== insert ====================== //
    iterator  insert_multi(const value_type& value);
    
    iterator  insert_multi(value_type&& value) {
        return emplace_multi(tinystl::move(value));
    }

    iterator  insert_multi(iterator hint, const value_type& value) {
        return emplace_multi_use_hint(hint, value);
    }

    iterator  insert_multi(iterator hint, value_type&& value) {
        return emplace_multi_use_hint(hint, tinystl::move(value));
    }

    template <class InputIterator>
    void      insert_multi(InputIterator first, InputIterator last) {
        size_type n = tinystl::distance(first, last);
        THROW_LENGTH_ERROR_IF(node_count_ > max_size() - n, "rb_tree<T, Comp>'s size too big");
        for (; n > 0; --n, ++first)
            insert_multi(end(), *first);
    }

    
    tinystl::pair<iterator, bool> insert_unique(const value_type& value);
    
    tinystl::pair<iterator, bool> insert_unique(value_type&& value) {
        return emplace_unique(tinystl::move(value));
    }

    iterator  insert_unique(iterator hint, const value_type& value) {
        return emplace_unique_use_hint(hint, value);
    }

    iterator  insert_unique(iterator hint, value_type&& value) {
        return emplace_unique_use_hint(hint, tinystl::move(value));
    }

    template <class InputIterator>
    void      insert_unique(InputIterator first, InputIterator last) {
        size_type n = tinystl::distance(first, last);
        THROW_LENGTH_ERROR_IF(node_count_ > max_size() - n, "rb_tree<T, Comp>'s size too big");
        for (; n > 0; --n, ++first)
            insert_unique(end(), *first);
    }

插入元素的辅助函数

get_Insert_pos

以下两个函数用于获取指定 key 在红黑树中的插入位置，分为 multi 与 unique 两类。

/// @brief 获取插入位置，键值允许重复，返回一个 pair，其中 first 为插入位置（父节点），second 表示是否在左侧插入
template <class T, class Compare>
tinystl::pair<typename rb_tree<T, Compare>::base_ptr, bool>
rb_tree<T, Compare>::get_insert_multi_pos(const key_type& key) {
    auto x = root();
    auto y = header_;
    bool add_to_left = true;
    while (x != nullptr) {
        y = x;
        add_to_left = key_comp_(key, value_traits::get_key(static_cast<node_ptr>(x)->value));
        x = add_to_left ? x->left : x->right;
    }
    return tinystl::make_pair(y, add_to_left);
}

/// @brief 获取插入位置，键值不允许重复，返回一个 pair<pair, bool>，其中 first 为一个 pair，
///       first.first 为插入位置，first.second 表示是否在左侧插入，second 表示是否插入成功
template <class T, class Compare>
tinystl::pair<tinystl::pair<typename rb_tree<T, Compare>::base_ptr, bool>, bool>
rb_tree<T, Compare>::get_insert_unique_pos(const key_type& key) {
    auto x = root();
    auto y = header_;
    bool add_to_left = true;  // 树为空时也在 header_ 左边插入
    while (x != nullptr) {
        y = x;
        add_to_left = key_comp_(key, value_traits::get_key(static_cast<node_ptr>(x)->value));
        x = add_to_left ? x->left : x->right;
    }
    iterator it(y);  // y 为插入点的父节点
    // iterator it = iterator(y);
    if (add_to_left) {
        // 如果树为空树或插入点在最左节点处，肯定可以插入新的节点 
        if (it == begin()) {
            return tinystl::make_pair(tinystl::make_pair(y, true), true);
        }
        // 否则（插入节点的父节点不是最左节点）
        else --it;
    }

    // 表明新节点没有重复
    if (key_comp_(value_traits::get_key(*it), key)) {
        return tinystl::make_pair(tinystl::make_pair(y, add_to_left), true);
    }
    // 进行至此，表示新节点与现有节点键值重复
    return tinystl::make_pair(tinystl::make_pair(y, add_to_left), false);
}

insert_node_at

以下的函数用于在指定的位置插入节点，值得注意的是，由于插入节点很有可能改变 root、leftmost、rightmost，因此这两个函数始终要维护正确的信息。此外，该函数会调用前面分析过的 rb_tree_insert_rebalance() 来维持红黑树的平衡。

/// @brief 在 x 位置插入 node 节点，add_to_left 表示是否在左侧插入
template <class T, class Compare>
typename rb_tree<T, Compare>::iterator
rb_tree<T, Compare>::insert_node_at(base_ptr x, node_ptr node, bool add_to_left) {
    node->parent = x;
    // auto base_node = node->get_base_ptr();
    auto base_node = static_cast<base_ptr>(node);
    if (x == header_) {
        root() = base_node;
        leftmost() = base_node;
        rightmost() = base_node;
    }
    else if (add_to_left) {
        x->left = base_node;
        if (x == leftmost()) leftmost() = base_node;
    }
    else {
        x->right = base_node;
        if (x == rightmost()) rightmost() = base_node;
    }
    rb_tree_insert_rebalance(base_node, root());
    ++node_count_;
    return iterator(node);
}

/// @brief 在 x 位置插入节点，节点值为 value，add_to_left 表示是否在左侧插入
template <class T, class Compare>
typename rb_tree<T, Compare>::iterator
rb_tree<T, Compare>::insert_value_at(base_ptr x, const value_type& value, bool add_to_left) {
    node_ptr node = create_node(value);
    node->parent = x;
    // auto base_node = node->get_base_ptr();
    auto base_node = static_cast<base_ptr>(node);
    if (x == header_) {
        root() = base_node;
        leftmost() = base_node;
        rightmost() = base_node;
    }
    else if (add_to_left) {
        x->left = base_node;
        if (x == leftmost()) leftmost() = base_node;
    }
    else {
        x->right = base_node;
        if (x == rightmost()) rightmost() = base_node;
    }
    rb_tree_insert_rebalance(base_node, root());
    ++node_count_;
    return iterator(node);
}

插入元素

有了上面两个辅助函数的帮助，插入元素的过程就是依次调用 get_insert_pos() 与 insert_node_at()。同时，红黑树中也有 emplace 与 insert 两种不同的插入方式。

/// @brief 就地插入元素，键值允许重复
template <class T, class Compare>
template <class ...Args>
typename rb_tree<T, Compare>::iterator
rb_tree<T, Compare>::emplace_multi(Args&& ...args) {
    THROW_LENGTH_ERROR_IF(node_count_ > max_size() - 1, "rb_tree<T, Compare>'s size too big");
    node_ptr node = create_node(tinystl::forward<Args>(args)...);
    auto res = get_insert_multi_pos(value_traits::get_key(node->value));
    return insert_node_at(res.first, node, res.second);
}

/// @brief 就地插入元素，键值不允许重复
/// @return 返回一个 pair，其中 first 为插入位置，second 表示是否插入成功
template <class T, class Compare>
template <class ...Args>
tinystl::pair<typename rb_tree<T, Compare>::iterator, bool>
rb_tree<T, Compare>::emplace_unique(Args&& ...args) {
    THROW_LENGTH_ERROR_IF(node_count_ > max_size() - 1, "rb_tree<T, Compare>'s size too big");
    node_ptr node = create_node(tinystl::forward<Args>(args)...);
    auto res = get_insert_unique_pos(value_traits::get_key(node->value));

    // 插入成功
    if (res.second) {
        return tinystl::make_pair(insert_node_at(res.first.first, node, res.first.second), true);
    }
    destroy_node(node);
    return tinystl::make_pair(iterator(res.first.first), false);
}

/// @brief 插入元素，节点键值允许重复
template <class T, class Compare>
typename rb_tree<T, Compare>::iterator
rb_tree<T, Compare>::insert_multi(const value_type& value) {
    THROW_LENGTH_ERROR_IF(node_count_ > max_size() - 1, "rb_tree<T, Compare>'s size too big");
    auto res = get_insert_multi_pos(value_traits::get_key(value));
    return insert_value_at(res.first, value, res.second);
}

/// @brief 插入元素，节点键值不允许重复，返回一个 pair，若插入成功，pair 的第二参数为 true，否则为 false
template <class T, class Compare>
tinystl::pair<typename rb_tree<T, Compare>::iterator, bool>
rb_tree<T, Compare>::insert_unique(const value_type& value) {
    THROW_LENGTH_ERROR_IF(node_count_ > max_size() - 1, "rb_tree<T, Compare>'s size too big");
    auto res = get_insert_unique_pos(value_traits::get_key(value));
    // 插入成功
    if (res.second) {
        return tinystl::make_pair(insert_value_at(res.first.first, value, res.first.second), true);
    }
    return tinystl::make_pair(res.first.first, false);
}

rb-tree 的元素删除

删除操作相较于插入操作会简单些，但有一部分删除操作需要用到 find() 函数，会在后面讲解。

// ====================== erase ====================== //

iterator  erase(iterator hint);

size_type erase_multi(const key_type& key);
size_type erase_unique(const key_type& key);

void erase(iterator first, iterator last);

void clear();

以下的删除节点函数基本都会调用 erase()，在 erase() 中可以看出，就是调用了 rb_tree_erase_rebalance() 函数，再销毁被要被删除的节点，再维护好 node_count_ 的信息。

/// @brief 删除 hint 位置的节点
/// @return 返回删除节点的后继节点
template <class T, class Compare>
typename rb_tree<T, Compare>::iterator
rb_tree<T, Compare>::erase(iterator hint) {
    // auto node = hint.node->get_node_ptr();
    auto node = static_cast<node_ptr>(hint.node);
    iterator next(node);
    ++next;
    // 重新平衡
    rb_tree_erase_rebalance(hint.node, root(), leftmost(), rightmost());
    destroy_node(node);
    --node_count_;
    return next;
}

/// @brief 删除键值等于 key 的元素，返回删除的个数
template <class T, class Compare>
typename rb_tree<T, Compare>::size_type
rb_tree<T, Compare>::erase_multi(const key_type& key) {
    auto res = equal_range_multi(key);
    auto count = tinystl::distance(res.first, res.second);
    erase(res.first, res.second);
    return count;
}

/// @brief 删除键值等于 key 的元素，返回删除的个数
template <class T, class Compare>
typename rb_tree<T, Compare>::size_type
rb_tree<T, Compare>::erase_unique(const key_type& key) {
    auto it = find(key);
    if (it == end()) return 0;
    erase(it);
    return 1;
}

/// @brief 删除 [first, last) 范围内的元素
template <class T, class Compare>
void rb_tree<T, Compare>::erase(iterator first, iterator last) {
    if (first == begin() && last == end()) clear();
    else {
        while (first != last) erase(first++);
    }
}

/// @brief 清空 rb-tree
template <class T, class Compare>
void rb_tree<T, Compare>::clear() {
    if (node_count_ != 0) {
        erase_since(root());
        leftmost() = header_;
        root() = nullptr;
        rightmost() = header_;
        node_count_ = 0;
    }
}

值得注意的是 clear() 函数，内部调用了一个名为 erase_since() 的函数，这个函数通过递归调用的方式销毁掉以 x 为根节点的所有节点。

/// @brief 从 x 开始递归删除树
template <class T, class Compare>
void rb_tree<T, Compare>::erase_since(base_ptr x) {
    while (x != nullptr) {
        erase_since(x->right);
        auto y = x->left;
        // destroy_node(x->get_node_ptr());
        destroy_node(static_cast<node_ptr>(x));
        x = y;
    }
}

rb-tree 的元素查找

由于红黑树也是一棵二叉搜索树，因此查找指定元素可以以 logn 的复杂度快速完成，除了 find() 外，rb-tree 还提供了 count、euqal_range 等查询接口，这部分的原理十分简单，就不再赘述。

public:  // 查找相关操作
    iterator              find(const key_type& key);

    size_type             count_multi(const key_type& key) const {
        auto p = equal_range_multi(key);
        return static_cast<size_type>(tinystl::distance(p.first, p.second));
    }

    size_type             count_unique(const key_type& key) const {
        return find(key) == end() ? 0 : 1;
    }

    iterator              lower_bound(const key_type& key);
    iterator              upper_bound(const key_type& key);

    tinystl::pair<iterator, iterator>
    equal_range_multi(const key_type& key) {
        return tinystl::pair<iterator, iterator>(lower_bound(key), upper_bound(key));
    }

    tinystl::pair<iterator, iterator>
    equal_range_unique(const key_type& key) {
        iterator it = find(key);
        auto next = it;
        return it == end() ? tinystl::make_pair(it, it) : tinystl::make_pair(it, ++next);
    }

/// @brief 查找键值为 key 的元素，若找到，返回指向该元素的迭代器，否则返回 end()
template <class T, class Compare>
typename rb_tree<T, Compare>::iterator
rb_tree<T, Compare>::find(const key_type& key) {
    auto y = header_;  // 最后一个不小于 key 的节点
    auto x = root();   // 当前节点
    while (x != nullptr) {
        // key 小于等于 x 键值，向左走
        if (!key_comp_(value_traits::get_key(static_cast<node_ptr>(x)->value), key)) {
            y = x;
            x = x->left;
        }
        else x = x->right;
    }
    iterator res(y);
    // 若 res 不等于 end()，且 key 小于等于 res 键值，返回 res
    if (res != end() && !key_comp_(key, value_traits::get_key(*res))) return res;
    return end();
}

/// @brief 键值不小于 key 的第一个位置
template <class T, class Compare>
typename rb_tree<T, Compare>::iterator
rb_tree<T, Compare>::lower_bound(const key_type& key) {
    auto y = header_;  // 最后一个不小于 key 的节点
    auto x = root();   // 当前节点
    while (x != nullptr) {
        // key 小于等于 x 键值，向左走
        if (!key_comp_(value_traits::get_key(static_cast<node_ptr>(x)->value), key)) {
            y = x;
            x = x->left;
        }
        else x = x->right;
    }
    return iterator(y);
}

/// @brief 键值不小于 key 的最后一个位置
template <class T, class Compare>
typename rb_tree<T, Compare>::iterator
rb_tree<T, Compare>::upper_bound(const key_type& key) {
    auto y = header_;  // 最后一个不小于 key 的节点
    auto x = root();   // 当前节点
    while (x != nullptr) {
        // key 小于 x 键值，向左走
        if (key_comp_(key, value_traits::get_key(static_cast<node_ptr>(x)->value))) {
            y = x;
            x = x->left;
        }
        else x = x->right;
    }
    return iterator(y);
}

rb-tree 的复制

关于红黑树的复制操作，也是通过递归调用的方式完成，如果出现了任何错误就销毁整棵复制的树。

/// @brief 复制一棵树，节点从 x 开始，p 为 x 的父节点
template <class T, class Compare>
typename rb_tree<T, Compare>::base_ptr
rb_tree<T, Compare>::copy_from(base_ptr x, base_ptr p) {
    auto top = clone_node(x);
    top->parent = p;
    try {
        // 若有右子树，则递归复制
        if (x->right) top->right = copy_from(x->right, top);
        p = top;
        x = x->left;
        while (x != nullptr) {
            auto y = clone_node(x);
            p->left = y;
            y->parent = p;
            if (x->right) y->right = copy_from(x->right, y);
            p = y;
            x = x->left;
        }
    }
    catch (...) {
        erase_since(top);
        throw;
    }
    return top;
}