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;
}