STL源码剖析（十三）关联式容器之rb_tree

STL源码剖析（十三）关联式容器之rb_tree

rt_tree是红黑树，在STL中，它并没有对使用者开放，但它非常重要，它作为set和map的底部支持

这篇文章不详细讨论红黑树的算法，着重讲解STL中红黑树的数据结构，如果需要了解红黑树，可以阅读 数据结构与算法之树（四）红黑树

一、rb_tree的数据结构

template <class Key, class Value, class KeyOfValue, class Compare,
          class Alloc = alloc>
class rb_tree {
  ...
proteced:
  typedef __rb_tree_node<Value> rb_tree_node;

  ...
public:
  typedef rb_tree_node* link_type;
  ...
protected:
  size_type node_count; // keeps track of size of tree
  link_type header;
  Compare key_compare;
  ...
};

rb_tree 中有三个成员变量，node_count 表示红黑树中的节点个数，header 表示头节点（这是STL红黑树的一个特性），key_compare 是一个仿函数，用于比较键值的大小，别忘了红黑树是一种二叉搜索树，是有序的，所以需要有一个比较节点的方法

rb_tree 的节点定义如下

struct __rb_tree_node_base
{
  typedef __rb_tree_color_type color_type;
  typedef __rb_tree_node_base* base_ptr;

  color_type color; 
  base_ptr parent;
  base_ptr left;
  base_ptr right;
}

template <class Value>
struct __rb_tree_node : public __rb_tree_node_base
{
  typedef __rb_tree_node<Value>* link_type;
  Value value_field;
};

每个节点有颜色 color，有父节点 parent，有左孩子 left，有右孩子 right，还有对应的值 vlur_field

rb_tree的红黑树是经过改造的，跟普通的红黑树还是有点区别，主要是多了 header节点，需要注意的是 header 不是根节点，是一个头节点，如下图所示

header 节点和根节点互为父节点，header 节点的左子节点指向红黑树最左边的节点，header 节点的右子节点指向红黑树最右边的节点

二、rb_tree的迭代器

struct __rb_tree_base_iterator
{
    
}

template <class Value, class Ref, class Ptr>
struct __rb_tree_iterator : public __rb_tree_base_iterator
{

}

__rb_tree_iterator 继承于 __rb_tree_base_iterator，先查看__rb_tree_base_iterator，其定义如下

struct __rb_tree_base_iterator
{
  typedef __rb_tree_node_base::base_ptr base_ptr;
  typedef bidirectional_iterator_tag iterator_category;
  typedef ptrdiff_t difference_type;
  base_ptr node;

  void increment();
  void decrement();
}

其中包含一个成员变量 node，其对应的类型为 typedef __rb_tree_node_base* base_ptr;

node 指向红黑树中的一个节点，其中的 increment 是跳到红黑树下一个节点的操作，对于二叉搜索树，使用中序遍历即可以按照顺序遍历整棵树，可想 increment 实现的就是非递归的中序遍历，如下

void increment()
{
  if (node->right != 0) { //跳到右子树最左的节点
    node = node->right;

    /* 跳到最左的节点 */
    while (node->left != 0)
      node = node->left;
  }
  else {
    /* 跳到第一个不为右节点的节点 */
    base_ptr y = node->parent;
    while (node == y->right) {
      node = y;
      y = y->parent;
    }

    /* 跳到父节点 */
    if (node->right != y)
      node = y;
  }
}

decrement 是返回上一个节点，其代码如下

void decrement()
{
  if (node->color == __rb_tree_red &&
      node->parent->parent == node)
    node = node->right;
  else if (node->left != 0) {
    base_ptr y = node->left;
    while (y->right != 0)
      y = y->right;
    node = y;
  }
  else {
    base_ptr y = node->parent;
    while (node == y->left) {
      node = y;
      y = y->parent;
    }
    node = y;
  }
}

下面再来看看 __rb_tree_iterator 的定义

template <class Value, class Ref, class Ptr>
struct __rb_tree_iterator : public __rb_tree_base_iterator
{
  /* STL迭代器的设计规范 */
  typedef Value value_type;
  typedef Ref reference;
  typedef Ptr pointer;
  typedef __rb_tree_iterator<Value, Value&, Value*>             iterator;

  reference operator*() const { return link_type(node)->value_field; }

  self& operator++() { increment(); return *this; }
  self& operator--() { decrement(); return *this; }
};

它并没有增加过多的内容，只是重载的自增和自减还有取值操作符，而它的实现也是基于基类的 increment 和 decrement 操作

三、rb_tree的操作

3.1 构造函数

默认构造

rb_tree(const Compare& comp = Compare())
  : node_count(0), key_compare(comp) { init(); }

首先初始化节点数，然后初始化键值比较的仿函数，然后调用 init 初始化 header 节点

下面是 init 的定义

void init() {
  header = get_node(); //分配内存
  color(header) = __rb_tree_red; //将节点设置为红色
  
  root() = 0; //将header的parent设置为NULL

  /* 将 header 的左右孩子设置为本身 */
  leftmost() = header;
  rightmost() = header;
}

其中的 root 函数定义如下

link_type& root() const { return (link_type&) header->parent; }

3.2 析构函数

~rb_tree() {
  clear();
  put_node(header);
}

首先调用 clear 清除所有的节点，然后调用 put_node 释放掉 header 节点的内存

clear 的定义如下

void clear() {
   if (node_count != 0) {
     __erase(root()); //销毁整棵树
       
     /* 初始化 header 节点 */
     leftmost() = header;
     root() = 0;
     rightmost() = header;
     node_count = 0;
}

首先通过 __erase 销毁整棵树，然后再是初始化 header 节点

__erase 传递根节点，通过递归的方式来清除整颗树

template <class Key, class Value, class KeyOfValue, class Compare, class Alloc>
void rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::__erase(link_type x) {
                                // erase without rebalancing
  while (x != 0) {
    __erase(right(x)); //递归到最右节点
    link_type y = left(x); //获取左子树
    destroy_node(x); //清除当前节点
    x = y; //跳到左子树继续递归
  }
}

3.3 插入元素

红黑树的插入方法支持 insert，不支持 push，因为红黑树是有序的，插入的节点会自动排序

insert 有两种方法，一种是 insert_equal 支持键值重复，一种是 insert_unique ，键值必须独一无二

红黑树的插入操作较为复杂，这里就不展开讨论了，只需要记住插入的接口就行

 iterator insert_unique(iterator position, const value_type& x);
iterator insert_equal(iterator position, const value_type& x);

3.4 删除元素

删除的操作是 erase，其接口如下

void erase(iterator position); //指定位置
size_type erase(const key_type& x); //指定键值

3.5 其他操作

begin

返回红黑树指向第一个元素的迭代器

link_type& leftmost() const { return (link_type&) header->left; }

iterator begin() { return leftmost(); } //红黑树的最左边节点

end

返回红黑树指向结尾节点的迭代器

iterator end() { return header; } //返回头节点

find

查找算法，STL的算法中页有find函数，但是红黑树自己定义的find函数。使用STL全局的find算法也可以工作，但是由于红黑树自身的特点，提供的一个find成员函数，大大提高查找效率

template <class Key, class Value, class KeyOfValue, class Compare, class Alloc>
typename rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::iterator 
rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::find(const Key& k) {
  link_type y = header;        // Last node which is not less than k. 
  link_type x = root();        // Current node. 

  while (x != 0) 
    if (!key_compare(key(x), k))
      y = x, x = left(x);
    else
      x = right(x);

  iterator j = iterator(y);   
  return (j == end() || key_compare(k, key(j.node))) ? end() : j;
}

红黑树是一种二叉搜索树，所以查找的方法也就是二分法

**lower_bound **

找到第一个大于等于给定值的节点

template <class Key, class Value, class KeyOfValue, class Compare, class Alloc>
typename rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::iterator 
rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::lower_bound(const Key& k) {
  link_type y = header; /* Last node which is not less than k. */
  link_type x = root(); /* Current node. */

  while (x != 0) 
    if (!key_compare(key(x), k))
      y = x, x = left(x);
    else
      x = right(x);

  return iterator(y);
}

**upper_bound **

找到第一个大于给定值的节点

template <class Key, class Value, class KeyOfValue, class Compare, class Alloc>
typename rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::iterator 
rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::upper_bound(const Key& k) {
  link_type y = header; /* Last node which is greater than k. */
  link_type x = root(); /* Current node. */

   while (x != 0) 
     if (key_compare(k, key(x)))
       y = x, x = left(x);
     else
       x = right(x);

   return iterator(y);
}