C++ 使用avltree实现一个set/map

本文使用C++实现一个基于AVLTree实现的map/set,不同于网上大部分的实现,本文争取按照C++STL风格来编写。

前言

其实网上关于数据结构的文章已经满大街了,但是大多数的实现都非常简单,虽然有些使用Java/C#/Python版本的代码甚至可能和其标准库相差无几,但是C++不一样,C++的STL的代码复杂度看上去显然是高于这些实现的。所以笔者希望能够尽可能地按照STL的风格来实现AVLTree。这样做一是希望可以把数据结构和C++更好地结合起来,让这个玩具可以更加漂亮,二是希望可以通过这个例子分享一些C++中的小细节,毕竟C++这门语言的细节太多了,很多细节可能不掌握对工作也没有什么影响,但是细节貌似又广泛存在于STL甚至folly/abseil这些库里面。

需要注意的是,和笔者之前的文章一样,重点不会停留在数据结构部分,我们假设读者对AVLTree已经有了一定程度的了解并且拥有实现其的能力。

AvlTree结点

为了方便遍历,我们使用三个指针,分别代表该节点的父节点,左孩子和右孩子。

struct avl_node
{
    avl_node* m_link[3];
};

头结点的设计

一个好的头结点的设计可能会让我们的算法变得简洁。考虑到我们可能希望可以快速获取二叉树的最值以及根结点,我们可以让头结点的左右孩子分别指向最大最小值,那么最后的父节点自然就指向了根结点。喜欢阅读源码的读者肯定看过类似的设计:

struct avl_node_base
{
    avl_node_base* father;
    avl_node_base* lchild;
    avl_node_base* rchild;
    int height;
};

template <typename T>
struct avl_node : avl_node_base
{
    T value; 
};

template <typename T>
struct avl_tree
{
    // 对于头结点来说,value并非必要的,因为我们的操作实际上都是从根结点开始的。
    avl_node_base header; 
};

不难发现大多数的BalancedBinarySearchTree的操作都是类似的,如AVLTree,RedBlackTree,毕竟它们都是基于BinarySearchTree的,所以我们先将一些基本的操作提取出来。我们采用模板来实现,C++23引入了deducing this。

struct node_interface
{
    template <typename Node>
    Node* parent(this Node& self)
    {
        return self.m_link[0];
    }

    template <typename Node>
    Node* lchild(this Node& self)
    {
        return self.m_link[1];
    }

    template <typename Node>
    Node* rchild(this Node& self)
    {
        return self.m_link[2];
    }

    template <typename Node>
    void parent(this Node& self, std::type_identity_t<Node>* node)
    {
        self.m_link[0] = node;
    }

    template <typename Node>
    void lchild(this Node& self, std::type_identity_t<Node>* node)
    {
        self.m_link[1] = node;
    }

    template <typename Node>
    void rchild(this Node& self, std::type_identity_t<Node>* node)
    {
        self.m_link[2] = node;
    }
};

为了保证下面代码可以正常编译,我们需要使用std::type_identity_t。

//[1]
template <typename Node>
void right(this Node& self, Node* node);

// [2]
template <typename Node>
void right(this Node& self, std::type_identity_t<Node>* node);

using node = avl_node<int>;
node* n1;
n1->right(nullptr);

// 对于[1]来说,第一个Node推导为 node, 第二个Node推导为 std::nullptr_t,两个Node类型不一致。
// 对于[2]来说,第一个Node推导为 node,第二个Node不参与推导,此时参数类型为 node*

基本操作

有了基本的set/get方法之后,我们来实现一些基本的BinarySearchTreeNode操作。

template <typename Node>
constexpr Node* maximum(this Node& self)
{
    auto x = std::addressof(self);
    for (; x->rchild(); x = x->rchild());
    return x;
}

template <typename Node>
constexpr Node* minimum(this Node& self)
{
    auto x = std::addressof(self);
    for (; x->lchild(); x = x->lchild());
    return x;
}

template <typename Node>
constexpr Node* decrement(this Node& self)
{
    auto x = std::addressof(self);

    if (x->lchild())
    {
        return x->lchild()->maximum();
    }
    else
    {
        auto y = x->parent();

        // 当前结点的设计可以保证这里的y一定不为nullptr
        while (x == y->lchild())
        {
            x = y;
            y = y->parent();
        }

        if (x->lchild() != y)
        {
            x = y;
        }

        return x;
    }
}

template <typename Node>
constexpr Node* increment(this Node& self)
{
    auto x = std::addressof(self);

    if (x->rchild())
    {
        return x->rchild()->minimum();
    }
    else
    {
        auto y = x->parent();

        while (x == y->rchild())
        {
            x = y;
            y = y->parent();
        }

        if (x->rchild() != y)
        {
            x = y;
        }

        return x;
    }
}

increment/decrement的实现会随着设计的不同发生改变,我们让头结点一开始lchildrchild指向自己。

header

对于一些特殊结点,我们有:

bool is_leaf(const Node* node)
{
    return !node->lchild() and !node->rchild();
}

bool is_root(const Node* node)
{
    // 我们让根结点的父结点指向header
    return node->parent() == header;
}

bool is_header(const Node* node)
{
    // 我们可以给根结点设置一个特殊的高度
    return node->m_height == -1;
}

BalancedBinarySearchTree一般是通过自旋操作来维持某种平衡的,接下来我们将实现基本的左旋和右旋操作。

/*
*     x            y              
*       \   =>   /    
*         y    x
*/
template <typename Node>    
void rotate_left(this Node& self, Node*& root)
{
    auto x = std::addressof(self);
    auto y = x->rchild();

    x->rchild(y->lchild());

    if (y->lchild())
    {
        y->lchild()->parent(x);
    }

    y->parent(x->parent());

    // 这里的x->parent一定不为空
    if (x == root)
    {
        root = y;
    }
    else if (x == x->parent()->lchild()) 
    {
        x->parent()->lchild(y);
    }
    else
    {
        x->parent()->rchild(y);
    }
    
    y->lchild(x);
    x->parent(y);
}

/*
*     x           y
*    /     =>      \
*   y               x
*/
template <typename Node>
void rotate_right(this Node& self, Node*& root)
{
    auto x = std::addressof(self);
    auto y = x->lchild();

    x->lchild(y->rchild());
    
    if (y->rchild())
    {
        y->rchild()->parent(x);
    }

    y->parent(x->parent());

    // 这里的x->parent一定不为空
    if (x == root)
    {
        root = y;
    }
    else if (x == x->parent()->rchild())
    {
        x->parent()->rchild(y);
    }
    else
    {
        x->parent()->lchild(y);
    }
    
    y->rchild(x);
    x->parent(y);
}

我们的设计可以减少很多对指针判空的情况,不过具体的算法实现我们在这里并不会过多赘述,否则本文的篇幅会过长,而且它们不是本文的重点。

接下来是AVLTree特有的操作,代码参考了avlmini,读者如果想了解细节可以参阅Github,当然该作者在知乎上也做了解释。

void avl_tree_fix_l(avl_node* header)
{
    auto x = this;
    auto r = x->rchild();
    int lh0 = height(r->lchild());
    int rh0 = height(r->rchild());

    if (lh0 > rh0)
    {
        r->rotate_right(header->parent());
        r->update_height();
        r->parent()->update_height();
    }

    x->rotate_left(header->parent());
    x->update_height();
    x->parent()->update_height();
}

void avl_tree_fix_r(avl_node* header)
{
    auto x = this;
    auto l = x->lchild();
    int lh0 = height(l->lchild());
    int rh0 = height(l->rchild());

    if (lh0 < rh0)
    {
        l->rotate_left(header->parent());
        l->update_height();
        l->parent()->update_height();
    }

    x->rotate_right(header->parent());
    x->update_height();
    x->parent()->update_height();
}

static int height(const avl_node* node)
{
    return node ? node->m_height : 0;
}

void update_height()
{
    int lh = height(lchild());
    int rh = height(rchild());
    m_height = std::max(lh, rh) + 1;
}

当某个结点的左子树和右子树的高度差大于2的时候,AVLTree会通过自旋操作来平衡高度。

平衡操作只会发生在插入和删除操作上,查找并不会改变树的高度,所以我们额外提供两个函数,一个用于插入操作,一个用于删除操作。

// 将当前结点插入到AVLTree中。
// 当insert_left为True时,将当前结点为p的左孩子,否则为右孩子。
// 考虑到平衡操作可能会让头节点操作发生变化,我们将头结点也作为参数传入。
void insert_and_rebalance(bool insert_left, avl_node* p, avl_node& header);

// 将当前节点从AVLTree中删除。
// 考虑到平衡操作可能会让头节点操作发生变化,我们将头结点也作为参数传入。
avl_node* rebalance_for_erase(avl_node& header);

经常阅读源码的读者肯定对上述接口很熟悉。

void insert_and_rebalance(bool insert_left, avl_node* p, avl_node& header)
{
    auto x = this;

    x->parent(p);
    x->lchild(nullptr);
    x->rchild(nullptr);
    x->m_height = 1;

    if (insert_left)
    {
        p->lchild(x);

        if (p == &header)
        {
            header.parent(x);
            header.rchild(x);
        }
        else if (p == header.lchild())
        {
            header.lchild(x);
        }
    }
    else
    {
        p->rchild(x);

        if (p == header.rchild())
        {
            header.rchild(x);
        }
    }

    // rebalance
    x->avl_tree_rebalance_insert(&header);
}

这个操作并不是很复杂,我们首先将结点插入,然后维护一下头结点,最后尝试平衡此树。细心的读者肯定会发现,分支语句中的代码并不是对称的,这是因为我们默认对一个空树插入结点时,insert_left为True。

void avl_tree_rebalance_insert(avl_node* header)
{
    auto x = this;

    for (x = x->parent(); x != header; x = x->parent())
    {
        int lh = height(x->lchild());
        int rh = height(x->rchild());
        int h = std::max(lh, rh) + 1;

        if (height(x) == h) 
        {
            break;
        }
        
        x->m_height = h;

        int diff = lh - rh;

        if (diff <= -2)
        {
            x->avl_tree_fix_l(header);
        }
        else if (diff >= 2)
        {
            x->avl_tree_fix_r(header);
        }
    }
}

具体如何平衡,并不是重点,想了解细节的读者可以在知乎上找到相关文章。删除操作就接口而言并没有需要特别解释的地方,代码之后会给出,这里直接省略。

至此为止,结点部分已经全部解释了,到目前为止,所有的操作均不涉及Compare部分,只是一些确定的操作,具体每个函数的参数该传入何值,由AVLTree决定。

迭代器的设计

我们直接将树的结点和迭代器作为树的内部类,因为它们可能会随着部分模板参数的变化而变化。这不是一个最好的方法,因为那些对其无关紧要的模板参数发生改变会让编译器多生成一份代码,但是比较简洁。

template <typename KeyValue,
    typename Compare,
    typename Allocator,
    bool UniqueKey, typename Node>
class tree;

解释一下参数:

  • KeyValue 用于切换Set/Map
  • Compare 用于比较两个元素
  • Allocator 用于内存分配
  • UniqueKey 用于切换Map/MultiMap,Set/MultiSet
  • Node 结点的种类

需要注意的是,本文的UniqueKey参数始终为True,笔者并不打算实现MultiSet/MultiMap。

对于KeyValue,我们提供两个辅助类:

template <typename T>
struct identity
{
    using key_type = T;

    using value_type = T;

    template <typename U>
    static constexpr auto&& operator()(U&& x)
    {
        return (U&&)x;
    }
};

template <typename T1, typename T2>
struct select1st
{
    using key_type = T1;

    using value_type = std::pair<const T1, T2>;

    template <typename U>
    static constexpr auto&& operator()(U&& x)
    {
        return ((U&&)x).first;
    }
};

由于Set和Map拥有不同的value_type,而实际上我们是通过比较key_type来确定元素插入的位置,所以我们需要一个辅助类来帮助我们获取元素的key。

接下来我们简单完善一下的代码:

template <typename KeyValue,
    typename Compare,
    typename Allocator,
    bool UniqueKey, typename Node>
class tree
{
public:

    using key_value = KeyValue;
    using value_type = typename KeyValue::value_type;
    using key_type = typename KeyValue::key_type;
    using size_type = std::size_t;
    using allocator_type = Allocator;
    using difference_type = std::ptrdiff_t;
    using key_compare = Compare;
    using value_compare = Compare;
    using reference = value_type&;
    using const_reference = const value_type&;
    using pointer = std::allocator_traits<Allocator>::pointer;
    using const_pointer = std::allocator_traits<Allocator>::const_pointer;

    struct tree_node : public Node
    {
        value_type m_val;

        value_type* value_ptr()
        {
            return std::addressof(m_val);
        }

        const value_type* value_ptr() const
        {
            return std::addressof(m_val);
        }

        Node* base()
        {
            return static_cast<Node*>(this);
        }

        const Node* base() const
        {
            return static_cast<const Node*>(this);
        }
    };

    using iterator = tree_iterator;
    // 注意如果我们的iterator本身就具有const的属性,那么std::const_iterator<iterator>推导的结果还是iterator
    using const_iterator = std::const_iterator<iterator>;
    using reverse_iterator = std::reverse_iterator<iterator>;
    using const_reverse_iterator = std::reverse_iterator<const_iterator>;
};

当我们能够确定所有类型之后,我们就可以完善我们迭代器的代码了。

struct tree_iterator 
{
    using link_type = Node*;
    using value_type = value_type;
    using difference_type = std::ptrdiff_t;
    using iterator_category = std::bidirectional_iterator_tag;

    // C++20之后无需定义pointer
    // using pointer = ...

    // 对于set<int>来说,key_type和value_type是一样的,我们可以通过这个特点来区分set/map。
    using reference = std::conditional_t<std::is_same_v<key_type, value_type>, const value_type&, value_type&>;

    // 我们在迭代器内存储一个结点的指针
    link_type m_ptr;

    constexpr tree_iterator() = default;
    constexpr tree_iterator(const tree_iterator&) = default;

    constexpr tree_iterator(link_type ptr) : m_ptr(ptr) { }

    // 以下接口不是必要的,这里只是为了方便调试,迭代器只需要重载相关操作符即可
    // 这里的this会使得iterator在调用这些函数的时候复制自身传入值
    // 迭代器只包含一个指针,结构很简单,和string_view类似,这里不必要传入const&
    constexpr link_type link(this tree_iterator it)
    {
        return it.m_ptr;
    }

    constexpr tree_iterator up(this tree_iterator it)
    {
        return tree_iterator(it.m_ptr->parent());
    }

    constexpr tree_iterator left(this tree_iterator it) 
    {    
        return tree_iterator(it.m_ptr->lchild());
    }

    constexpr tree_iterator right(this tree_iterator it)
    {
        return tree_iterator(it.m_ptr->rchild());
    }

    // 头结点自增我们让他指向leftmost
    constexpr tree_iterator& operator++()
    {
        m_ptr = m_ptr->is_header() ? m_ptr->lchild() : m_ptr->increment();
        return *this;
    }

    // 头结点自减我们让他指向rightmost
    constexpr tree_iterator& operator--()
    {
        m_ptr = m_ptr->is_header() ? m_ptr->rchild() : m_ptr->decrement();
        return *this;
    }

    constexpr tree_iterator operator++(int)
    {
        tree_iterator tmp = *this;
        ++*this;
        return tmp;
    }

    constexpr tree_iterator operator--(int)
    {
        tree_iterator tmp = *this;
        --*this;
        return tmp;
    }

    constexpr reference operator*(this tree_iterator it) 
    {
        return *(static_cast<tree_node*>(it.m_ptr)->value_ptr());
    }

    friend constexpr bool operator==(tree_iterator lhs, tree_iterator rhs) 
    {
        return lhs.m_ptr == rhs.m_ptr;
    }

    // C++20之后!=操作可以由==操作自动生成
};

相关的操作我们已经在avl_node中提供了,迭代器只需要简单封装一下即可。

树的设计

到目前为止我们已经介绍了结点和迭代器部分的设计,最后就是树本身了。

template <typename KeyValue,
    typename Compare,
    typename Allocator,
    bool UniqueKey, typename Node>
class tree
{
    using node_allocator = typename std::allocator_traits<Allocator>::template rebind_alloc<tree_node>;
    using node_alloc_traits = std::allocator_traits<node_allocator>;
    
    [[no_unique_address]] Compare m_cmp;
    [[no_unique_address]] Allocator m_alloc;
    Node m_header;
    size_type m_size;
};

搜索树的操作都是Key基于查找的,所以我们先提供一些辅助函数帮助我们获取结点的Key:

static const key_type& keys(const tree_node* node)
{
    return KeyValue()(*node->value_ptr());
}

static const key_type& keys(const node_base* node)
{
    return keys(static_cast<const tree_node*>(node));
}

插入操作

// 我们这里采用模板的形式而非const key_type&,
// 因为当我们的Compare定义了transparent的时候,
// 我们并不要求传入的类型必须是key_type
template <typename K>
std::pair<node_base*, node_base*> get_insert_unique_pos(const K& k)
{
    auto y = &m_header, x = header()->parent();
    bool comp = true;

    while (x)
    {
        y = x;
        comp = m_cmp(k, keys(x));
        // 当comp为True时,说明当前的k的值小于结点的键值,应该向左孩子方向遍历
        x = comp ? x->lchild() : x->rchild();
    }

    // y始终指向x的父结点
    iterator j{ y };

    if (comp)
    {
        if (j == begin())
        {
            // 如果comp为True且j为leftmost,那么k比树中所有结点的键值都小
            return { x, y };
        }
        else
        {
            // 否则我们不能确定是否存在一个键值和k相等的结点
            --j;
        }
    }

    // 此时满足j <= k

    // 再次比较,如果j < k则说明可以插入
    if (m_cmp(KeyValue()(*j), k))
    {
        return { x, y };
    }

    // 否则键值已经存在
    return { j.m_ptr, nullptr };
}

这种写法一般不会出现在网络上常见的教程中,所以笔者在上述的代码中添加了注释。我们通过下面两张图来理解一下。
img
img

当我们找到需要插入的位置的时候,我们就尝试对接点进行插入操作:

template <typename Arg>
std::pair<iterator, bool> insert_unique(Arg&& arg)
{
    auto [x, p] = get_insert_unique_pos(KeyValue()(arg));

    if (p)
    {
        return { insert_value(x, p, (Arg&&)arg), true };
    }
    
    return { x, false };
}

template<typename Arg>
iterator insert_value(node_base* x, node_base* p, Arg&& arg)
{
    bool insert_left = (x != 0 || p == &m_header
        || m_cmp(KeyValue()(arg), keys(p)));

    auto z = create_node((Arg&&)arg);

    z->insert_and_rebalance(insert_left, p, m_header);
    ++m_size;
    return iterator(z);
}

熟悉STL数据结构的读者肯定会发现除了insert以外,还会有一些函数带有emplace字眼,比如

class vector
{
    void push_back(const T& x)
    { emplace_back(x); }

    template <typename... Args>
    reference emplace_back(Args&&... args);
};

同理,我们这里依旧可以如此操作:

std::pair<iterator, bool> insert(const value_type& value)
{
    return emplace(value);
}

std::pair<iterator, bool> insert(value_type&& value)
{
    return emplace(std::move(value));
}

template <typename... Args>
std::pair<iterator, bool> emplace(Args&&... args)
{
    if constexpr (detail::emplace_helper<value_type, Args...>::value)
    {
        return insert_unique((Args&&)args...);
    }
    else
    {
        // 我们的插入操作是基于key的比较的,所以我们必须先使用args构造值
        // value_handle在构造函数中使用allocator构造元素
        // 在析构函数中使用allocator销毁元素
        value_handle<value_type, allocator> handle(m_alloc, (Args&&)args...);
        return insert_unique(*handle);
    }
}

考虑到搜索树和普通的vector不同的是,搜索树需要一个完成的元素来进行key的比较才可以确定最终的位置,所以我们必须先把值构造出来在进行插入,在这里我们实际上就简单地在栈上面构造一个元素然后move一下。在这里我们一开始做了一个简单的判断,这是一个小的优化,如果我们传入的参数只有一个并且其类型正好是const value&/value_type&&,直接调用insert_unique。

到此为止我们最核心的插入操作其实已经完成了,接下来大部分的insert函数我们只需要复用此操作即可完成。

C++17引入了一些特殊的insert方法:

insert_return_type insert(node_type&& nh)
node_type extract(const_iterator position)
node_type extract(const key_type& x)
template <typename K> node_type extract(K&& x)

这些方法我们直接复用容器的结点,减少内存分配的次数。

我们在这里可以了解到node_type是如何设计的。
我们按照要求简单复现一下:

template <typename NodeAllocator>
struct node_handle_base
{
    using allocator_type = NodeAllocator;
    using ator_traits  = std::allocator_traits<allocator_type>;
    using pointer = typename ator_traits::pointer;
    using container_node_type = typename std::pointer_traits<pointer>::element_type;

    pointer m_ptr = nullptr; // exposition-only

    // 实际上我们并不需要使用std::optional,因为其内部有一个bool变量
    // 用于表示是否有值,但是对于我们这种情况来说。我们可以利用m_ptr来代替bool变量,
    // 当指针为空时,结点无效,allocator自然也就是无效的,反之有效。
    // 不过本文就直接使用std::optional。
    std::optional<allocator_type> m_alloc; // exposition-only
};

template <typename Iterator, typename NodeHandle>
struct node_insert_return
{
    Iterator position;
    bool inserted;
    NodeHandle node;
};

很显然,node_handle里面可以存储一个结点的指针以及一个allocator,后者用于销毁元素并释放内存。

对于extract函数,我们只需要找到结点然后将结点从树中删除即可。

node_type extract(const_iterator position)
{
    return extract_node(position.base().m_ptr);
}

node_type extract(const key_type& x)
{
    return extract(find((x)));
}

template <typename K>
    requires detail::transparent<Compare>
node_type extract(K&& x)
{
    return extract(find((x)));
}

node_type extract_node(node_base* node)
{
    if (node == std::addressof(m_header))
    {
        return node_type(nullptr, m_alloc);
    }
    else
    {
        // unlink and reset the node   
        node->rebalance_for_erase(m_header);
        node->parent(nullptr);
        node->lchild(nullptr);
        node->rchild(nullptr);
        node->init();
        --m_size;
        return node_type(static_cast<tree_node*>(node), m_alloc);
    }
}

结点的插入也是类似的,找到位置后直接插入即可。

insert_return_type insert(node_type&& nh)
{
    if (nh.empty())
    {
        return { end(), false, node_type(nullptr, m_alloc) };
    }

    tree_node* node = std::exchange(nh.m_ptr, nullptr);
    auto [x, p] = get_insert_unique_pos(*node->value_ptr());

    return p == nullptr ? 
        insert_return_type(iterator(x), false, node_type(node, m_alloc)) : 
        insert_return_type(insert_node(x, p, node->base()), true, node_type(nullptr, m_alloc));
}

iterator insert_node(node_base* x, node_base* p, node_base* z)
{
    bool insert_left = (x != 0 || p == &m_header || m_cmp(keys(z), keys(p)));

    z->insert_and_rebalance(insert_left, p, m_header);
    ++m_size;
    return iterator(z);
}

到这里插入操作部分就结束了。

查找操作

template <typename K>
iterator lower_bound_impl(const K& k)
{
    auto y = &m_header, x = header()->parent();

    while (x)
    {
        if (!m_cmp(keys(x), k))
        {
            y = x, x = x->lchild();
        }
        else
        {
            x = x->rchild();
        }
    }
    
    return { y };
}

这个操作核心部分和get_insert_unique_pos是一样的。

template <typename K = key_type>
iterator lower_bound(const key_arg_t<K>& x)
{
    return lower_bound_impl(x);
}

template <typename K = key_type>
const_iterator lower_bound(const key_arg_t<K>& x) const
{
    return const_cast<tree&>(*this).lower_bound(x);
}

STL的set中有四种形式的重载:

iterator lower_bound( const Key& key );  // (1)
	
const_iterator lower_bound( const Key& key ) const; // (2)
	
template< class K >
iterator lower_bound( const K& x ); // (3)

template< class K >
const_iterator lower_bound( const K& x ) const; // (4)

当Compare定义了transparent时(3)(4)参与重载。

我们通过一个模板的小技巧Non-deduced contexts可以把重载的数量下降到两个。

template <typename T>
concept has_transparent = requires { typename T::is_transparent; };

template <typename... Ts>
concept transparent = (has_transparent<Ts> && ...); 

template <bool IsTransparent>
struct key_arg_helper
{
    template <typename K1, typename K2>
    using type = K1;
};

template <>
struct key_arg_helper<false>
{
    template <typename K1, typename K2>
    using type = K2;
};

template <bool IsTransparent, typename K1, typename K2>
using key_arg = typename key_arg_helper<IsTransparent>::template type<K1, K2>;

template <typename U> 
using key_arg_t = detail::key_arg<detail::transparent<Compare>, U, key_type>;

我们尽可能通过lower_bound来实现其他查找函数,这样的好处是代码比较简洁,不容易出错。

template <typename K = key_type>
iterator find(const key_arg_t<K>& x)
{
    iterator lower = lower_bound(x);
    return (lower == end() || m_cmp(x, keys(lower.m_ptr))) 
            ? end() : lower;
}

template <typename K = key_type>
const_iterator find(const key_arg_t<K>& x) const
{
    return const_cast<tree&>(*this).find(x);
}

template <typename K = key_type>
std::pair<iterator, iterator> equal_range(const key_arg_t<K>& x)
{
    iterator lower = lower_bound(x);
    iterator upper = (lower == end() || m_cmp(x, keys(lower.m_ptr))) ? lower : std::next(lower); 
    return std::make_pair(lower, upper);
}

template <typename K = key_type>
std::pair<const_iterator, const_iterator> equal_range(const key_arg_t<K>& x) const
{
    return const_cast<tree&>(*this).equal_range(x);
}

template <typename K = key_type>
iterator upper_bound(const key_arg_t<K>& x)
{
    return equal_range(x).second;
}

template <typename K = key_type>
const_iterator upper_bound(const key_arg_t<K>& x) const
{
    return const_cast<tree&>(*this).upper_bound(x);
}

template <typename K = key_type>
bool contains(const key_arg_t<K>& x) const
{
    return find(x) != end();
}

template <typename K = key_type>
size_type count(const key_arg_t<K>& x) const
{
    return contains(x);
}

删除操作

对于删除,我们只需要找到对应的结点将其从树中移除然后释放内存即可。

void erase_by_node(node_base* x)
{      
    x->rebalance_for_erase(m_header);
    drop_node(static_cast<tree_node*>(x));
    m_size--;
}

template <typename K>
size_type erase_by_key(const K& x)
{
    iterator iter = find(x);

    if (iter != end())
    {
        erase_by_node(iter.m_ptr);
        return 1;
    }

    return 0;
}

到此为止,树的操作以及全部结束。对于一个带有allocator的容器来说,move/copy/swap方法也是有很多需要注意的细节,笔者也写过类似的文章,这里不再赘述。

扩展到set/map

我们可以通过继承的方式来进行扩展,因为大多数方法都完全一样,如果组合的话代码会显得冗余。

template <typename K, typename V, typename Compare, typename Allocator, bool Unique, typename Node>
class tree_map : public tree<select1st<K, V>, Compare, Allocator, Unique, Node>,
                 public unique_associative_container_indexer_interface
{
    using base = tree<select1st<K, V>, Compare, Allocator, Unique, Node>;
public:

    using mapped_type = V;
    using typename base::value_type;
    using typename base::iterator;
    using typename base::const_iterator;
};

对于map来说多了一个value_compare,我们只需要按照要求去实现它即可。

struct value_compare 
{
    bool operator()(const value_type& lhs, const value_type& rhs) const
    {
        return m_comp(lhs.first, rhs.first);
    }

protected:
    friend class tree_map;
    
    value_compare(Compare compare) : m_comp(compare) { }

    Compare m_comp;
};

细心的读者会发现,许多函数Cppref上都给出了一个等价形式,也许这个形式并非最高校的,但是其正确性是可以得到保证的。

V& operator[](const K &key)
{
    return this->try_emplace(key).first->second;
}

V& operator[](K &&key)
{
    return this->try_emplace(std::move(key)).first->second;
}

template <typename... Args>
std::pair<iterator, bool> try_emplace(const K &k, Args &&...args)
{
    return try_emplace_impl(k, (Args &&)args...);
}

template <typename... Args>
std::pair<iterator, bool> try_emplace(K &&k, Args &&...args)
{
    return try_emplace_impl(std::move(k), (Args &&)args...);
}

template <typename... Args>
std::pair<iterator, bool> try_emplace(const_iterator, const K &k, Args &&...args)
{
    return try_emplace_impl(k, (Args &&)args...);
}

template <typename... Args>
std::pair<iterator, bool> try_emplace(const_iterator, K &&k, Args &&...args)
{
    return try_emplace_impl(std::move(k), (Args &&)args...);
}

template <typename M>
std::pair<iterator, bool> insert_or_assign(const K &k, M &&obj)
{
    return insert_or_assign_impl(k, (M &&)obj);
}

template <typename M>
std::pair<iterator, bool> insert_or_assign(K &&k, M &&obj)
{
    return insert_or_assign_impl(std::move(k), (M &&)obj);
}

protected:

template <typename KK, typename M>
std::pair<iterator, bool> insert_or_assign_impl(KK&& k, M&& obj)
{
    auto [x, p] = this->get_insert_unique_pos(k);
    if (p) 
    {
        auto z = this->create_node((KK&&)k, (M&&)obj);
        return { this->insert_node(x, p, z), true };
    }
    auto j = iterator(x);
    *j = (M&&)obj;
    return { j, false };
}

template <typename KK, typename... Args>
std::pair<iterator, bool> try_emplace_impl(KK&& k, Args&&... args)
{
    auto [x, p] = this->get_insert_unique_pos(k);
    if (p) 
    {
        auto z = this->create_node(
            std::piecewise_construct, 
            std::forward_as_tuple((KK&)k), 
            std::forward_as_tuple((Args&&)args...));
        return { this->insert_node(x, p, z), true };
    }
    return { x, false };
}

扩展到MultiSet/MultiMap

这就留给读者自行思考了。需要注意的是MultiSet/MultiMap的接口可能和Set/Map不同。

需要源码的读者可以点击这里

附录

关于头结点设计的一些问题

有些人喜欢这样设计:

struct avl_node
{
    avl_node* father;
    avl_node* lchild;
    avl_node* rchild;
    int height;
    T value;
};

struct avl_tree
{
    avl_node* header;
};

如果我们放一根指针在里面,那么当我们初始化AVLTree的时候就需要考虑是否需要为header申请一片内存。

  • 申请内存保证header不为null。显然申请内存是需要付出代价的,我们可能并不希望一个普通的构造函数会耗费过多的资源。

  • 不申请内存,实际操作的时候再考虑是否需要申请内存。这样我们的代码中会出现大量的if判断,这样不仅会让我们的代码更加冗余,甚至更容易出错。

这样看来不使用指针是一个好主意,那如果采用值呢?

struct avl_node
{
    avl_node* father;
    avl_node* lchild;
    avl_node* rchild;
    int height;
    // T value;
    alignas(T) unsigned char buffer[sizeof(T)];
};

struct avl_tree
{
    avl_node header;
};

这种做法需要修改value的表达形式,否则对于没有默认构造函数的类型会产生编译错误。由于头结点并不包含值,所以本文也未采用这种方法。

关于迭代器设计的一些问题

在这里我们只想谈论一下++/--操作。在C++中,显然我们不应该去解引用一个end()迭代器,我们不知道这个无效的迭代器里面到底是什么,有时候里面可能是一个空指针,比如某些非循环链表的设计。

// 这些代码只是例子,不会出现在最终代码中,所以我们采用不同的命名以区分。
struct LinkListNode
{
    LinkListNode* next = nullptr;
    int value;
};

struct LinkListIterator
{
    LinkListNode* node;

    auto& operator*() const { return node->value; }
};

struct LinkList
{
    LinkListIterator end() { return LinkListIterator(nullptr); }
};

但是对于++/--这样的操作呢?如果我们对一个空指针解引用获取next字段,那么程序可能就崩溃了。我们这种树的设计不会造成这个问题,但是可能会带来一些其他问题:

std::set<int>().begin()++; // infinity loop

MSVC通过引入了一个新的字段来进行判断来避免这种错误代码导致的死循环。我们也希望可以通过一些小的改变来避免这些情况,所以我们设计的迭代器是循环的。

幸运的是我们并不需要引入新的字段,因为对于AVLTree来说,所有的结点高度都是一个非负数,所以我们可以将头结点(header/end/sentinel)的高度设置为-1,这样就可以快速判断一个结点是否为头结点。我们不必担心特殊高度对我们的平衡操作带来影响,因为我们平衡操作的终止条件为:

for (... ; node != header; ...) { ... }

这并不会涉及到header的高度。实际上RedBlackTree也不需要,我们可以引入新的颜色:

enum struct Color 
{
    Red, Black, Sentinel
};

struct RedBlackNode
{
    RedBlackNode* parent;
    RedBlackNode* lchild;
    RedBlackNode* rchild;
    Color color;
};

我们只需要把头结点的颜色设置为Sentinel即可。实际的操作和AVLTree一样并不会涉及到头结点。

综上所诉,作者希望可以通过将其设计为循环结构来避免++/--造成的一些问题。

关于C++中的Compare一些讨论

在C++中,Compare一般返回的是一个bool值,这一点上和其他语言可能有些区别。在其他语言中,两个对象比较可能返回一个整数,比如下面这种形式:

SomeObject obj1;
SomeObject obj2;

int res = obj1.CompareTo(obj2);

if (res < 0)
{
    // maybe res == -1
    Console::WriteLine("obj1 is less than obj2");
}
else if (obj > 0)
{
    // maybe res == 1
    Console::WriteLine("obj1 is greater than obj2");
}
else
{
    Console::WriteLine("obj1 is equal/equivalent to obj2");
}

C++20引入了<=>运算符,不知道未来是否可能使用,但是就已经存在的组件而言,笔者认为不太会改变。有些教材可能会采用类似于下面的写法:

Compare compare;

if (compare(x, y))
{
    // x is less than y
}
else if (compare(y, x))
{
    // x is greater than y
}
else
{
    // x is equivalent to y
}

这个写法乍一看多了一个分支语句,同时多了一次compare的调用,我们简单讨论一下这个写法。

我们编写一个简单的BinarySearch,采用如上形式:

template <typename R, typename T, typename Compare = std::less<T>>
bool BinarySearch(R&& r, T value, Compare compare)
{
    auto left = std::begin(r);
    auto right = std::end(r);
    
    while (left < right)
    {
        auto middle = left + (right - left) / 2;
        if (compare(value, *middle))
        {
            right = middle;
        }
        else if (compare(*middle, value))
        {
            left = middle + 1;
        }
        else
        {
            return true;
        }
    }

    return false;
}

接下来我们定义两个Compare,内含计数器,帮助我们统计比较的次数。

struct CounterCompare1
{
    template <typename L, typename R>
    static constexpr bool operator()(L&& l, R&& r)
    {
        ++count;
        return l < r;    
    } 

    inline static int count = 0;
};

struct CounterCompare2
{
    template <typename L, typename R>
    static constexpr bool operator()(L&& l, R&& r)
    {
        ++count;
        return l < r;    
    } 

    inline static int count = 0;
};

最后我们随机产生一些测试样例帮助我们测试一下:

std::random_device rd;

constexpr int N = 100'000;
constexpr int M = N * 1000;

std::vector<int> GetRandomVector(int count)
{
    std::vector<int> v;
    for (int i = 0; i < count; ++i)
    {
        v.push_back(rd() % M);
    }
    std::sort(std::begin(v), std::end(v));
    return v;
}

int main(int argc, char const *argv[])
{
    auto source = GetRandomVector(N);

    auto search_values = GetRandomVector(N);

    for (auto value : search_values)
    {
        auto b1 = std::ranges::binary_search(source, value, CounterCompare1{});
        auto b2 = BinarySearch(source, value, CounterCompare2{});
    
        if (b1 != b2)
        {
            std::cout << "Error: " << value << std::endl;
        }
    }

    std::cout << "BinarySearch1: " << CounterCompare1::count << std::endl;      // 1769074
    std::cout << "BinarySearch2: " << CounterCompare2::count << std::endl;      // 2484674

    return 0;
}

我们发现即使改变N和M的值,两次compare写法的比较次数总是会高于STL,而STL使用的就是和我们搜索树一样的一次compare写法。

关于emplace的优化讨论

emplace一般都是接受若干个参数:

template <typename... Args>
std::pair<iterator, bool> emplace(Args&&... args);

对于搜索树来说我们可能需要先使用args...将值构建出来然后通过比较才能确定最终元素插入的位置。但是有时候,我们传入的参数本身类型也是value_type,在这种情况下,构建元素显得并没有过多的意义,所以我们对参数的数量和类型简单判断一下:

template <typename T, typename... Args>
struct emplace_helper
{
    static constexpr bool value = []() 
    {
        if constexpr (sizeof...(Args) != 1)
            return false;
        else
            return std::conjunction_v<std::is_same<T, std::remove_cvref_t<Args>>...>;
    }();
};

template <typename... Args>
std::pair<iterator, bool> emplace(Args&&... args)
{
    if constexpr (detail::emplace_helper<value_type, Args...>::value)
    {
        return insert_unique((Args&&)args...);
    }
    else
    {
        // we use a value_handle to help us destroy 
        value_handle<value_type, allocator_type> handle(m_alloc, (Args&&)args...);
        return insert_unique(*handle);
    }
}

我们这里并没有直接构造一个结点然后尝试插入,对于MultiSet来说,这是一个很好的做法,因为MultiSet的插入必然会成功,但是对于Set来说,如果当前容器中已经有此元素,插入会失败。插入失败会让我们多申请和释放一次内存,所以我们直接在栈上面构造一个值然后尝试move。

我们的写法都是以UniqueKey参数为True为前提写的,实际上如果为False的话,许多函数的写法都需要做相应的调整,不过并不会调整太多。我们可以使用if-constexpr或者继承的方式进行修改。

扩展到RedBlackTree

我们的树可以通过传入不同的结点变成不同种类的树,对于SplayTree这种查找也会修改自身的平衡树我们可能无法满足,不过这种树也不满足STL对set容器的要求,不必考虑。

接下来我们将stdlibc++中红黑树结点代码复制过来并简单修改来适配我们的代码:

struct _Rb_tree_node_base : node_interface
{
    // Different from standard red black tree, we add another color sentinel, 
    // for each node except header, the color is red or black. And in this way, 
    // the iterator can be a circle. Since compiler will fill some bytes for 
    // struct, it will always occupy sizeof(_Rb_tree_node_base*) bytes. 
    enum _Rb_tree_color
    {
        _S_red,
        _S_black,
        _S_sentinel
    };

    _Rb_tree_node_base* _M_parent;
    _Rb_tree_node_base* _M_left;
    _Rb_tree_node_base* _M_right;
    _Rb_tree_color _M_color;

    _Rb_tree_node_base* lchild() { return _M_left; }
    const _Rb_tree_node_base* lchild() const { return _M_left; }
    void lchild(_Rb_tree_node_base* node) { _M_left = node; }

    _Rb_tree_node_base* rchild() { return _M_right; }
    const _Rb_tree_node_base* rchild() const { return _M_right; }
    void rchild(_Rb_tree_node_base* node) { _M_right = node; }

    _Rb_tree_node_base* parent() { return _M_parent; }
    const _Rb_tree_node_base* parent() const { return _M_parent; }
    void parent(_Rb_tree_node_base* node) { _M_parent = node; }

    std::string to_string() const
    {
        switch (_M_color)
        {
        case _Rb_tree_color::_S_red: return "R";
        case _Rb_tree_color::_S_black: return "B";
        default: return "S";
        }
    }

    static constexpr void reset(_Rb_tree_node_base* node)
    {
        // node->_M_parent = nullptr;
        // node->_M_left = node->_M_right = node;
        node->_M_color  = _S_sentinel;
    }

    void as_empty_tree_header()
    {
        _Rb_tree_node_base::reset(this);
    }

    static constexpr void init(_Rb_tree_node_base* node)
    {
        // node->_M_left = node->_M_right = node->_M_parent = nullptr;
        node->_M_color  = _S_red;
    }

    void init()
    {
        _Rb_tree_node_base::init(this);
    }

    static constexpr bool is_header(const _Rb_tree_node_base *node)
    {
        return node->_M_color  == _S_sentinel;
    }

    bool is_header() const
    {
        return _Rb_tree_node_base::is_header(this);
    }

    static constexpr void clone(_Rb_tree_node_base *x, const _Rb_tree_node_base *y)
    {
        x->_M_color = y->_M_color ;
    }

    void clone(const _Rb_tree_node_base* y)
    {
        _Rb_tree_node_base::clone(this, y);
    }

    void insert_and_rebalance(bool insert_left, _Rb_tree_node_base* p, _Rb_tree_node_base& header)
    {
        _Rb_tree_node_base::insert_and_rebalance(
            insert_left,
            this, 
            p, 
            header
        );
    }

    static constexpr void
    insert_and_rebalance(const bool __insert_left,
                                    _Rb_tree_node_base *__x,
                                    _Rb_tree_node_base *__p,
                                    _Rb_tree_node_base &__header) 
    {
        __header._M_color  = _S_red;
        _Rb_tree_node_base *&__root = __header._M_parent;

        // Initialize fields in new node to insert.
        __x->_M_parent = __p;
        __x->_M_left = 0;
        __x->_M_right = 0;
        __x->_M_color  = _S_red;

        // Insert.
        // Make new node child of parent and maintain root, leftmost and
        // rightmost nodes.
        // N.B. First node is always inserted left.
        if (__insert_left)
        {
            __p->_M_left = __x; // also makes leftmost = __x when __p == &__header

            if (__p == &__header)
            {
                __header._M_parent = __x;
                __header._M_right = __x;
            }
            else if (__p == __header._M_left)
                __header._M_left = __x; // maintain leftmost pointing to min node
        }
        else
        {
            __p->_M_right = __x;

            if (__p == __header._M_right)
                __header._M_right = __x; // maintain rightmost pointing to max node
        }
        // Rebalance.
        while (__x != __root && __x->_M_parent->_M_color == _S_red)
        {
            _Rb_tree_node_base *const __xpp = __x->_M_parent->_M_parent;

            if (__x->_M_parent == __xpp->_M_left)
            {
                _Rb_tree_node_base *const __y = __xpp->_M_right;
                if (__y && __y->_M_color  == _S_red)
                {
                    __x->_M_parent->_M_color  = _S_black;
                    __y->_M_color  = _S_black;
                    __xpp->_M_color  = _S_red;
                    __x = __xpp;
                }
                else
                {
                    if (__x == __x->_M_parent->_M_right)
                    {
                        __x = __x->_M_parent;
                        tree_rotate_left(__x, __root);
                    }
                    __x->_M_parent->_M_color  = _S_black;
                    __xpp->_M_color  = _S_red;
                    tree_rotate_right(__xpp, __root);
                }
            }
            else
            {
                _Rb_tree_node_base *const __y = __xpp->_M_left;
                if (__y && __y->_M_color  == _S_red)
                {
                    __x->_M_parent->_M_color  = _S_black;
                    __y->_M_color  = _S_black;
                    __xpp->_M_color  = _S_red;
                    __x = __xpp;
                }
                else
                {
                    if (__x == __x->_M_parent->_M_left)
                    {
                        __x = __x->_M_parent;
                        tree_rotate_right(__x, __root);
                    }
                    __x->_M_parent->_M_color  = _S_black;
                    __xpp->_M_color  = _S_red;
                    tree_rotate_left(__xpp, __root);
                }
            }
        }
        __root->_M_color  = _S_black;
        __header._M_color  = _S_sentinel; // Make header black.
    }

    _Rb_tree_node_base* rebalance_for_erase(_Rb_tree_node_base& header)
    {
        return _Rb_tree_node_base::rebalance_for_erase(this, header);
    }

    static constexpr _Rb_tree_node_base* maximum(_Rb_tree_node_base* x)
    {
        assert(x && "x should not be nullptr");
        for (; x->_M_right; x = x->_M_right);
        return x;
    }

    _Rb_tree_node_base* minimum() { return _Rb_tree_node_base::minimum(this); }
    const _Rb_tree_node_base* minimum() const { return _Rb_tree_node_base::minimum(this); }

    _Rb_tree_node_base* maximum() { return _Rb_tree_node_base::maximum(this); }
    const _Rb_tree_node_base* maximum() const { return _Rb_tree_node_base::maximum(this); }

    static constexpr _Rb_tree_node_base* minimum(_Rb_tree_node_base* x)
    {
        assert(x && "x should not be nullptr");
        for (; x->_M_left; x = x->_M_left);
        return x;
    }

    static constexpr const _Rb_tree_node_base* maximum(const _Rb_tree_node_base* x)
    { return maximum(const_cast<_Rb_tree_node_base*>(x)); }

    static constexpr const _Rb_tree_node_base* minimum(const _Rb_tree_node_base* x)
    { return minimum(const_cast<_Rb_tree_node_base*>(x)); }

    /*
    *     x            y              
    *       \   =>   /    
    *         y    x
    */
    static constexpr void tree_rotate_left(_Rb_tree_node_base* x, _Rb_tree_node_base*& root)
    {
        _Rb_tree_node_base* y = x->_M_right;

        x->_M_right = y->_M_left;
        if (y->_M_left != 0)
            y->_M_left->_M_parent = x;
        y->_M_parent = x->_M_parent;

        // x->parent will never be nullptr, since header->parent == root and root->parent == header
        if (x == root)
            root = y;
        else if (x == x->_M_parent->_M_left) 
            x->_M_parent->_M_left = y;
        else
            x->_M_parent->_M_right = y;
        y->_M_left = x;
        x->_M_parent = y;
    }


    /*
    *     x        y                   
    *    /     =>    \
    *   y              x
    */
    static constexpr void tree_rotate_right(_Rb_tree_node_base* x, _Rb_tree_node_base*& root)
    {
        _Rb_tree_node_base* y = x->_M_left;

        x->_M_left = y->_M_right;
        if (y->_M_right != 0)
            y->_M_right->_M_parent = x;
        y->_M_parent = x->_M_parent;

        if (x == root)
            root = y;
        else if (x == x->_M_parent->_M_right)
            x->_M_parent->_M_right = y;
        else
            x->_M_parent->_M_left = y;
        y->_M_right = x;
        x->_M_parent = y;
    }


    static constexpr _Rb_tree_node_base *
    rebalance_for_erase(_Rb_tree_node_base *const __z,
                                    _Rb_tree_node_base &__header) 
    {

        __header._M_color  = _S_red;

        _Rb_tree_node_base *&__root = __header._M_parent;
        _Rb_tree_node_base *&__leftmost = __header._M_left;
        _Rb_tree_node_base *&__rightmost = __header._M_right;
        _Rb_tree_node_base *__y = __z;
        _Rb_tree_node_base *__x = 0;
        _Rb_tree_node_base *__x_parent = 0;

        if (__y->_M_left == 0)       // __z has at most one non-null child. y == z.
            __x = __y->_M_right;     // __x might be null.
        else if (__y->_M_right == 0) // __z has exactly one non-null child. y == z.
            __x = __y->_M_left;      // __x is not null.
        else
        {
            // __z has two non-null children.  Set __y to
            __y = __y->_M_right; //   __z's successor.  __x might be null.
            while (__y->_M_left != 0)
                __y = __y->_M_left;
            __x = __y->_M_right;
        }
        if (__y != __z)
        {
            // relink y in place of z.  y is z's successor
            __z->_M_left->_M_parent = __y;
            __y->_M_left = __z->_M_left;
            if (__y != __z->_M_right)
            {
                __x_parent = __y->_M_parent;
                if (__x)
                    __x->_M_parent = __y->_M_parent;
                __y->_M_parent->_M_left = __x; // __y must be a child of _M_left
                __y->_M_right = __z->_M_right;
                __z->_M_right->_M_parent = __y;
            }
            else
                __x_parent = __y;
            if (__root == __z)
                __root = __y;
            else if (__z->_M_parent->_M_left == __z)
                __z->_M_parent->_M_left = __y;
            else
                __z->_M_parent->_M_right = __y;
            __y->_M_parent = __z->_M_parent;
            std::swap(__y->_M_color , __z->_M_color );
            __y = __z;
            // __y now points to node to be actually deleted
        }
        else
        { // __y == __z
            __x_parent = __y->_M_parent;
            if (__x)
                __x->_M_parent = __y->_M_parent;
            if (__root == __z)
                __root = __x;
            else if (__z->_M_parent->_M_left == __z)
                __z->_M_parent->_M_left = __x;
            else
                __z->_M_parent->_M_right = __x;
            if (__leftmost == __z)
            {
                if (__z->_M_right == 0) // __z->_M_left must be null also
                    __leftmost = __z->_M_parent;
                // makes __leftmost == _M_header if __z == __root
                else
                    __leftmost = _Rb_tree_node_base::minimum(__x);
            }
            if (__rightmost == __z)
            {
                if (__z->_M_left == 0) // __z->_M_right must be null also
                    __rightmost = __z->_M_parent;
                // makes __rightmost == _M_header if __z == __root
                else // __x == __z->_M_left
                    __rightmost = _Rb_tree_node_base::maximum(__x);
            }
        }
        if (__y->_M_color  != _S_red)
        {
            while (__x != __root && (__x == 0 || __x->_M_color  == _S_black))
                if (__x == __x_parent->_M_left)
                {
                    _Rb_tree_node_base *__w = __x_parent->_M_right;
                    if (__w->_M_color  == _S_red)
                    {
                        __w->_M_color  = _S_black;
                        __x_parent->_M_color  = _S_red;
                        tree_rotate_left(__x_parent, __root);
                        __w = __x_parent->_M_right;
                    }
                    if ((__w->_M_left == 0 ||
                            __w->_M_left->_M_color  == _S_black) &&
                        (__w->_M_right == 0 ||
                            __w->_M_right->_M_color  == _S_black))
                    {
                        __w->_M_color  = _S_red;
                        __x = __x_parent;
                        __x_parent = __x_parent->_M_parent;
                    }
                    else
                    {
                        if (__w->_M_right == 0 || __w->_M_right->_M_color  == _S_black)
                        {
                            __w->_M_left->_M_color  = _S_black;
                            __w->_M_color  = _S_red;
                            tree_rotate_right(__w, __root);
                            __w = __x_parent->_M_right;
                        }
                        __w->_M_color  = __x_parent->_M_color ;
                        __x_parent->_M_color  = _S_black;
                        if (__w->_M_right)
                            __w->_M_right->_M_color  = _S_black;
                        tree_rotate_left(__x_parent, __root);
                        break;
                    }
                }
                else
                {
                    // same as above, with _M_right <-> _M_left.
                    _Rb_tree_node_base *__w = __x_parent->_M_left;
                    if (__w->_M_color  == _S_red)
                    {
                        __w->_M_color  = _S_black;
                        __x_parent->_M_color  = _S_red;
                        tree_rotate_right(__x_parent, __root);
                        __w = __x_parent->_M_left;
                    }
                    if ((__w->_M_right == 0 ||
                            __w->_M_right->_M_color  == _S_black) &&
                        (__w->_M_left == 0 ||
                            __w->_M_left->_M_color  == _S_black))
                    {
                        __w->_M_color  = _S_red;
                        __x = __x_parent;
                        __x_parent = __x_parent->_M_parent;
                    }
                    else
                    {
                        if (__w->_M_left == 0 || __w->_M_left->_M_color  == _S_black)
                        {
                            __w->_M_right->_M_color  = _S_black;
                            __w->_M_color  = _S_red;
                            //   tree_rotate_left(__w, __root);
                            tree_rotate_left(__w, __root);
                            __w = __x_parent->_M_left;
                        }
                        __w->_M_color  = __x_parent->_M_color ;
                        __x_parent->_M_color  = _S_black;
                        if (__w->_M_left)
                            __w->_M_left->_M_color  = _S_black;
                        //   tree_rotate_right(__x_parent, __root);
                        tree_rotate_right(__x_parent, __root);
                        break;
                    }
                }
            if (__x)
                __x->_M_color  = _S_black;
        }
        
        __header._M_color  = _S_sentinel;     // Is unnecessary?       
        return __y;
    }

};

using RedBlackTree = tree<::identity<int>, std::less<int>, std::allocator<int>, true, _Rb_tree_node_base>;

posted @   鸿钧三清  Views(3)  Comments(0Edit  收藏  举报
相关博文:
阅读排行:
· 阿里最新开源QwQ-32B,效果媲美deepseek-r1满血版,部署成本又又又降低了!
· 开源Multi-agent AI智能体框架aevatar.ai,欢迎大家贡献代码
· Manus重磅发布:全球首款通用AI代理技术深度解析与实战指南
· 被坑几百块钱后,我竟然真的恢复了删除的微信聊天记录!
· AI技术革命,工作效率10个最佳AI工具
点击右上角即可分享
微信分享提示