STL浅析 RB-tree（红黑树）

RB-tree(红黑树)是平衡二叉搜索树。RB-tree满足二叉搜索树的规则之外，还遵循以下特性：

• 每个节点不是红色就是黑色。
• 根节点为黑色。
• 如果节点为红色，其子节点必须为黑色。
• 任意一个节点到到NULL（树尾端）的任何路径，所含之黑色节点数必须相同。

红黑树原理

因为红黑树比较复杂，特别是旋转和变色很难描述，篇幅也会很长，所以这里推荐一下这个视频：
红黑树原理及平衡二叉树旋转详解

RB-tree节点设计

为了有更大的弹性，节点的结构分成了两层。
__rb_tree_node_base记录了节点的颜色、父节点指针以及左右子节点指针，
__rb_tree_node继承了__rb_tree_node_base并记录节点的值。

typedef bool __rb_tree_color_type;
const __rb_tree_color_type __rb_tree_red = false;//红色为0
const __rb_tree_color_type __rb_tree_black = true;//黑色为1

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;//当前节点的右子节点指针
    
    // 寻找当前节点为根节点的最小节点，最左边为最小节点
    static base_ptr minimum(base_ptr x)
    {
        while (x->left != 0) x = x->left;
        return x;
    }
    
    // 寻找当前节点为根节点的最大节点，最右边为最大节点
    static base_ptr maximum(base_ptr x)
    {
        while (x->right != 0) x = x->right;
        return x;
    }
};

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

RB-tree迭代器

为了让RB-tree实现成为一个泛型容器，就需要将迭代器也设计成两层，与节点设计相互配合。
__rb_tree_base_iterator设计了节点的前进和后退操作，__rb_tree_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()
    {
        if (node->right != 0) {   // 有右节点，一直向右节点的左子节点移动，直到最后得到自增结果
            node = node->right;
            while (node->left != 0)
                node = node->left;
        }
        else {    // 没有右节点
            base_ptr y = node->parent;

			// 该节点是父节点的右子节点
			// 一直上溯右节点，直到node不是右子节点
            while (node == y->right) {
                node = y;
                y = y->parent;
            }
            
            // 上溯完成后，上溯后节点是父节点的左子节点，那么父节点就是自增结果
            if (node->right != y)
                node = y;
                
           // 状况4，当前迭代器为根节点，且无右节点时，根节点本身为自增结果(依赖于hrader节点的特殊实现)
        }
    }
    
    // 自减
    void decrement()
    {
        if (node->color == __rb_tree_red && node->parent->parent == node)
            // 状况1，节点为红色，且父节点的父节点为自己，自减结果为右子节点
            // 此情况发生于对end()迭代器自减，依赖于header节点的特殊实现，header节点的右子节点是整颗树的最大节点
            node = node->right;
        else if (node->left != 0) {  // 有左子节点，找到左子树中的最大值
            base_ptr y = node->left;
            while (y->right != 0)
                y = y->right;
            node = y;
        }
        else {   // 不是根节点，又没有左子节点，一直上溯左节点，直到node不是左子节点，得到自减结果
            base_ptr y = node->parent;
            while (node == y->left) {
                node = y;
                y = y->parent;
            }
            node = y;
        }
    }
};

// 继承了__rb_tree_base_iterator
template <class Value, class Ref, class Ptr>
struct __rb_tree_iterator : public __rb_tree_base_iterator
{
    typedef Value value_type;
    typedef Ref reference;
    typedef Ptr pointer;
    typedef __rb_tree_iterator<Value, Value&, Value*>             iterator;
    typedef __rb_tree_iterator<Value, const Value&, const Value*> const_iterator;
    typedef __rb_tree_iterator<Value, Ref, Ptr>                   self;
    typedef __rb_tree_node<Value>* link_type;
    
    __rb_tree_iterator() {}
    __rb_tree_iterator(link_type x) { node = x; }
    __rb_tree_iterator(const iterator& it) { node = it.node; }
    
    reference operator*() const { return link_type(node)->value_field; }
#ifndef __SGI_STL_NO_ARROW_OPERATOR
    pointer operator->() const { return &(operator*()); }
#endif /* __SGI_STL_NO_ARROW_OPERATOR */
    
    self& operator++() { increment(); return *this; }
    self operator++(int) {
        self tmp = *this;
        increment();
        return tmp;
    }
    
    self& operator--() { decrement(); return *this; }
    self operator--(int) {
        self tmp = *this;
        decrement();
        return tmp;
    }
};

RB-tree数据结构

一下为RB-tree的定义：

template <class Key, class Value, class KeyOfValue, class Compare, class Alloc = alloc>
class rb_tree {
protected:
    typedef void* void_pointer;
    typedef __rb_tree_node_base* base_ptr;
    typedef __rb_tree_node<Value> rb_tree_node;
    typedef simple_alloc<rb_tree_node, Alloc> rb_tree_node_allocator;  // 专属空间配置器，配置单位为一个节点的空间
    typedef __rb_tree_color_type color_type;
public:
    typedef Key key_type;
    typedef Value value_type;
    typedef value_type* pointer;
    typedef const value_type* const_pointer;
    typedef value_type& reference;
    typedef const value_type& const_reference;
    typedef rb_tree_node* link_type;
    typedef size_t size_type;
    typedef ptrdiff_t difference_type;
protected:
    link_type get_node() { return rb_tree_node_allocator::allocate(); }
    void put_node(link_type p) { rb_tree_node_allocator::deallocate(p); }

    // 配置空间并且构造
    link_type create_node(const value_type& x) {
        link_type tmp = get_node();
        __STL_TRY{
            construct(&tmp->value_field, x);//构造内容
        }
        __STL_UNWIND(put_node(tmp));
        return tmp;
    }

    // 复制节点的值和颜色
    link_type clone_node(link_type x) {
        link_type tmp = create_node(x->value_field);
        tmp->color = x->color;
        tmp->left = 0;
        tmp->right = 0;
        return tmp;
    }

    void destroy_node(link_type p) {
        destroy(&p->value_field);//析构内容
        put_node(p);//释放空间
    }
public:
    // 迭代器定义
    typedef __rb_tree_iterator<value_type, reference, pointer> iterator;
    typedef __rb_tree_iterator<value_type, const_reference, const_pointer> const_iterator;
private:
    iterator __insert(base_ptr x, base_ptr y, const value_type& v);
    link_type __copy(link_type x, link_type p);
    void __erase(link_type x);

    // 初始化时，生成一个头结点
    void init() {
        header = get_node();
        color(header) = __rb_tree_red;   // 头结点颜色为红色，区分于根节点
        root() = 0;
        leftmost() = header;
        rightmost() = header;
    }
...
}

RB-tree的构造

template <class Key, class Value, class KeyOfValue, class Compare, class Alloc = alloc>
class rb_tree {
protected:
    typedef __rb_tree_node<Value> rb_tree_node;
    typedef simple_alloc<rb_tree_node, Alloc> rb_tree_node_allocator;  // 专属配置器，配置单位为一个节点的空间
public:
    // 默认构造，使用init函数生成一个带header
    rb_tree(const Compare& comp = Compare())
        : node_count(0), key_compare(comp) {
        init();
    }

   // 拷贝另一个树进行构造
    rb_tree(const rb_tree<Key, Value, KeyOfValue, Compare, Alloc>& x)
        : node_count(0), key_compare(x.key_compare)
    {
        header = get_node();
        color(header) = __rb_tree_red;
        if (x.root() == 0) {
            root() = 0;
            leftmost() = header;
            rightmost() = header;
        }
        else {
            __STL_TRY{
                root() = __copy(x.root(), header);
            }
            __STL_UNWIND(put_node(header));
            leftmost() = minimum(root());
            rightmost() = maximum(root());
        }
        node_count = x.node_count;
    }
    ~rb_tree() {
        clear();
        put_node(header);
    }
}
...

RB-tree的元素操作

元素插入insert_equal
insert_equal允许插入键值一致的节点，返回值是指向新增节点的迭代器。

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>::insert_equal(const Value& v)
{
    link_type y = header;
    link_type x = root();

    // 比较key值大小，小向左，大向右
    while (x != 0) {
        y = x;
        x = key_compare(KeyOfValue()(v), key(x)) ? left(x) : right(x);
    }
    // x为新值插入点，y为父节点，v是插入值
    return __insert(x, y, v);
}

元素插入insert_unique
insert_unique不允许插入键值重复的节点，返回值是pair值，包括是否插入成功以及插入成功后的新增节点迭代器。

template <class Key, class Value, class KeyOfValue, class Compare, class Alloc>
pair<typename rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::iterator, bool>
rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::insert_unique(const Value& v)
{
    link_type y = header;
    link_type x = root();
    bool comp = true;
    // 比较key值大小，小向左，大向右
    while (x != 0) {
        y = x;
        comp = key_compare(KeyOfValue()(v), key(x));
        x = comp ? left(x) : right(x);
    }
    
    iterator j = iterator(y);   // 父节点迭代器
    if (comp)                     // comp为true，表示插入节点较小，插入左侧
        if (j == begin())       // 如果父节点是最左节点，x为新值插入点，y为父节点，v是插入值
            return pair<iterator,bool>(__insert(x, y, v), true);
        else  // 如果父节点不是最左节点，将迭代器向前移动
            --j;
            
     // 比较向前移动后的节点和插入节点的键值大小，新值较大则插入右侧
    if (key_compare(key(j.node), KeyOfValue()(v)))
        return pair<iterator,bool>(__insert(x, y, v), true);
        
    //  否则，不大不小就是重复，那么返回重复的节点并告知失败
    return pair<iterator,bool>(j, false);
}

插入操作__insert
以下为真正的执行插入节点操作。

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>::
__insert(base_ptr x_, base_ptr y_, const Value& v) {
    link_type x = (link_type) x_;
    link_type y = (link_type) y_;
    link_type z;
    
    // 当插入节点为头结点或插入键值小于父节点时
    // x!=0，插入点不为NULL为特殊情况在insert_unique的另一个版本中使用
    if (y == header || x != 0 || key_compare(KeyOfValue()(v), key(y))) {
        z = create_node(v);
        left(y) = z;        // 插入父节点的左侧，当父节点为header时等价于leftmost() = z; 
        if (y == header) {
            root() = z;
            rightmost() = z;
        }
        else if (y == leftmost())
            leftmost() = z;           // 父节点为最小值，直接插入左侧
    }
    else {  // 插入键值大于父节点的情况
        z = create_node(v);
        right(y) = z;
        if (y == rightmost())
            rightmost() = z;            // 父节点为最大值，直接插入右侧
    }

    // 为插入节点设置父节点等
    parent(z) = y;
    left(z) = 0;
    right(z) = 0;

	// 进行平衡(改变颜色，进行旋转)
    __rb_tree_rebalance(z, header->parent);
    ++node_count;
    return iterator(z);
}

调整树 __rb_tree_rebalance

inline void
__rb_tree_rebalance(__rb_tree_node_base* x, __rb_tree_node_base*& root)
{
    // 因为规则：任意一个节点到到NULL（树尾端）的任何路径，所含之黑色节点数必须相同
    // 所以插入节点默认为红色，保证这条规则必然符合
    x->color = __rb_tree_red;  

    // 父节点颜色为红，且当前节点不为根节点(违反规则：如果节点为红色，其子节点必须为黑色)
    while (x != root && x->parent->color == __rb_tree_red) {
        if (x->parent == x->parent->parent->left) {
            // 父节点是祖父节点的左子节点
            __rb_tree_node_base* y = x->parent->parent->right;
            if (y && y->color == __rb_tree_red) {   
                // 伯父节点存在且为红色(伯父节点必然和父节点同色)
                // 因为规则：如果节点为红色，其子节点必须为黑色，插入点位红色不符合要求
                // 所以此时需要改变父节点和伯父节点的颜色为黑，祖父节点为红
                x->parent->color = __rb_tree_black;
                y->color = __rb_tree_black;
                x->parent->parent->color = __rb_tree_red;
                
                // 同时以祖父节点为调整起点继续调整
                // 如果祖父节点的父节点为黑色按照循环条件就调整结束
                // 否则继续调整
                x = x->parent->parent;
                // 注意：上文中介绍插入情况时，有伯父节点还需要进行一次旋转，实际代码中并没有
            }
            else {
                // 伯父节点不存在，需要进行旋转保持平衡
                if (x == x->parent->right) {
                   // 插入节点为右节点，内侧插入，需要先左旋再右旋
                   // 以父节点为旋转点进行左旋
                    x = x->parent;
                    __rb_tree_rotate_left(x, root);
                }
                
                // 使父节点颜色为黑，祖父节点颜色为红，这样旋转之后颜色和深度都能保证平衡
                x->parent->color = __rb_tree_black;
                x->parent->parent->color = __rb_tree_red;
                
               // 以祖父节点为旋转点进行右旋
                __rb_tree_rotate_right(x->parent->parent, root);
            }
        }
        else {
           // 父节点是祖父节点的右节点，和以上的情况对称
            __rb_tree_node_base* y = x->parent->parent->left;
            if (y && y->color == __rb_tree_red) {
                x->parent->color = __rb_tree_black;
                y->color = __rb_tree_black;
                x->parent->parent->color = __rb_tree_red;
                x = x->parent->parent;
            }
            else {
                if (x == x->parent->left) {
                    x = x->parent;
                    __rb_tree_rotate_right(x, root);
                }
                x->parent->color = __rb_tree_black;
                x->parent->parent->color = __rb_tree_red;
                __rb_tree_rotate_left(x->parent->parent, root);
            }
        }
    }

    // 按照规则根节点始终为黑
    root->color = __rb_tree_black;
}

// 左旋转
inline void
__rb_tree_rotate_left(__rb_tree_node_base* x, __rb_tree_node_base*& root)
{
    __rb_tree_node_base* y = x->right;  // 旋转点的右子节点
    // 1、将y的左子节点设置成x的右子节点
    x->right = y->left; 
    if (y->left != 0)
        y->left->parent = x;
    // 2、将y的父节点设置成x的父节点
    y->parent = x->parent;

    // 3、将y替换到x的位置
    if (x == root)
        root = y;
    else if (x == x->parent->left)
        x->parent->left = y;
    else
        x->parent->right = y;
  
    // 4、将x设置成y的左子节点
    y->left = x;
    x->parent = y;
}

// 右旋转
inline void
__rb_tree_rotate_right(__rb_tree_node_base* x, __rb_tree_node_base*& root)
{
    __rb_tree_node_base* y = x->left;  // 选装点的左子节点
    // 1、将y的右子节点设置成x的左子节点
    x->left = y->right;
    if (y->right != 0)
        y->right->parent = x;
    // 2、将y的父节点设置成x的父节点
    y->parent = x->parent;

    // 3、将y替换到x的位置
    if (x == root)
        root = y;
    else if (x == x->parent->right)
        x->parent->right = y;
    else
        x->parent->left = y;

    // 4、将x设置成y的右子节点
    y->right = x;
    x->parent = y;
}

元素查找

查找函数提供两个版本，其中一个返回const迭代器。

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; 
    link_type x = root(); 

    // 节点不为空
    while (x != 0)
        // 比较查找值和当前值的大小，大向右，小或等于向左
        // y用来保存x节点，最后返回结果，因为小于是向左移动，找到之后会一直向右直到叶子节点为空
        if (!key_compare(key(x), k))
            y = x, x = left(x);                    
        else
            x = right(x); 
            
    // 获取查找结果y的迭代器
    // 判断y是否为end()也就是header节点，此时返回end()
    // 判断k是否大于j的值，如果是说明没有找到，返回end()
    iterator j = iterator(y);
    return (j == end() || key_compare(k, key(j.node))) ? end() : j;
}

template <class Key, class Value, class KeyOfValue, class Compare, class Alloc>
typename rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::const_iterator
rb_tree<Key, Value, KeyOfValue, Compare, Alloc>::find(const Key& k) const {
    link_type y = header; 
    link_type x = root();

    while (x != 0) {
        if (!key_compare(key(x), k))
            y = x, x = left(x);
        else
            x = right(x);
    }
    const_iterator j = const_iterator(y);
    return (j == end() || key_compare(k, key(j.node))) ? end() : j;
}