red_black_tree的一个实现(c/c++)

参考《STL源码剖析》, 博客:http://blog.csdn.net/spch2008/article/details/9338923

#ifndef MY_RB_TREE_H#define MY_RB_TREE_H#include<iostream>#include<cstddef>using namespace std;#define t_classes <typename Key,typename Value,typename KeyOfValue,typename Compare> #define classes <Key,Value,KeyOfValue,Compare> typedef bool rb_tree_color_type;const rb_tree_color_type rb_tree_red = false;const rb_tree_color_type rb_tree_black = true;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 != NULL) x = x -> left;        return x;    }    static base_ptr maximum(base_ptr x){        while(x -> right != NULL) x = x -> right;        return x;    }    base_ptr increment(){        base_ptr x = this;        if(x->right != NULL){            x = x->right;            while(x->left != NULL) x = x->left;        }else{            base_ptr y = x ->parent;            while(x == y->right){                x = y;                y = y->parent;            }            if(x->right != y) //y != header else return header                x = y;        }        return x;    }    base_ptr decrement(){        base_ptr x = this;        if(x->color == rb_tree_red && x->parent->parent == x){ //x == header            x = x->right;        }else if(x ->left != NULL){            base_ptr y = x->left;            while(y->right != NULL) y = y->right;            x = y;        }else{            base_ptr y = x->parent;            while(x == y->left){                x = y;                y = y->parent;            }            x = y;        }        return x;    }};template <class Value>struct __rb_tree_node: public rb_tree_node_base{    typedef  __rb_tree_node<Value>* link_type;    Value value_field;};//KeyOfValue is a function object (functor). In this rb_tree, the value_field contains key and value actrually. We can define the structure of the value_field but we should provide this function object to extract the key from the value_field//Compare is a function object. if a < b return true else falsetemplate t_classesclass rb_tree{protected:    typedef  void* void_pointer;    typedef rb_tree_node_base* base_ptr;    typedef __rb_tree_node<Value> rb_tree_node;    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 new rb_tree_node();          }       void put_node(link_type p){        delete p;    }    link_type create_node (const value_type x){        link_type node = get_node();        //construct(&node->value_field,x);        node->value_field = value_type(x);    }    link_type clone_node(link_type x){        link_type  tmp = creat_node(x->value_field);        tmp->color = x->color;        tmp->left = NULL;        tmp->right = NULL;        return tmp;    }//if you used alloc and free in c, it should destroy the object first then release the memory    void destroy_node(link_type p){        put_node(p);    }protected:    size_type node_count; //number of nodes    link_type header; //this is an extra node for convinience, its parent is the head,left is the left most and right child is the right most node    Compare key_compare; //a compare function object (functor)    link_type& root() const {return (link_type &) header->parent;   }    link_type& leftmost() const {return (link_type &) header->left; }    link_type& rightmost() const {return (link_type &) header->right;   }    static link_type& left(link_type x){ return (link_type&)(x->left);}    static link_type& right(link_type x){ return (link_type&)(x->right);}    static link_type& parent(link_type x){ return (link_type&)(x->parent);}    static reference value(link_type x){ return x->value_field;}    static const Key key(link_type x){ return KeyOfValue()(value(x));}    static color_type& color(link_type x){ return (color_type&)(x->color);}    static link_type minimum(link_type x){        return (link_type) rb_tree_node_base::minimum(x);    }    static link_type maximum(link_type x){        return (link_type) rb_tree_node_base::maximum(x);    }private:    link_type insert(base_ptr x,base_ptr y,const value_type & v);    void init(){        header = get_node();        color(header) = rb_tree_red;        root() = 0;        leftmost() = header;        rightmost() = header;    }    void del_tree(link_type &root);public:    rb_tree(const Compare& comp = Compare()):node_count(0),key_compare(comp){init();}    ~rb_tree(){        clear();        put_node(header);    }//  rb_tree<Key, Value, KeyOfValue, Compare> & operator= (const rb_tree<Key,Value,KeyOfValue,Compare>& x);    Compare key_comp()const{return key_compare;}    bool empty(){ return node_count == 0; }    size_type size() const {return node_count;}public:    pair<link_type,bool> insert_unique(const value_type& v);    link_type insert_equal(const value_type& x);    link_type find(const key_type& v);    void del(link_type& x);    void clear();   private:    //void rb_tree_rebalance(rb_tree_node_base* x,rb_tree_node_base*& root);};template t_classesvoid rb_tree classes ::del_tree(link_type& root){    if(root == NULL) return;    del_tree(left(root));    del_tree(right(root));    put_node(root);}template t_classesvoid rb_tree classes ::clear(){    if(node_count > 0){        del_tree(root());        node_count = 0;        root() = NULL;        leftmost() = header;        rightmost() = header;    }}//x is position to insert value v. y is the parent of the node of the position.template t_classestypename rb_tree classes ::link_typerb_tree classes ::insert(base_ptr _x,base_ptr _y,const value_type & v){    link_type x = (link_type) _x;    link_type y = (link_type) _y;    link_type z;    //tree empty or x not leaf or cur_key < key_y, insert to left child    if(y == header || x != 0 || key_compare(KeyOfValue()(v), key(y))){        z = create_node(v);        left(y) = z;        if(y == header){            root() = z;            rightmost() = z;        }else if(y == leftmost())            leftmost() = z;    }else{ //right child        z = create_node(v);        right(y) = z;        if(y == rightmost())            rightmost() = z;    }    parent(z) = y;    left(z) = right(z) = NULL;    rb_tree_rebalance(z,header->parent);    ++node_count;    return z;}template t_classespair<typename rb_tree classes ::link_type,bool>rb_tree classes ::insert_unique(const value_type& v){    link_type y = header;    link_type x = root();    bool comp = true;    while(x != 0){        y = x;        comp = key_compare(KeyOfValue()(v),key(x));        x = comp ? left(x) : right(x); //less than or not    }       link_type j = y;    //after while, y is the parent of x, and x is the position to insert    if(comp){ //insert to the left        if(j == leftmost())            return pair<link_type,bool>(insert(x,y,v),true);            else            j = (link_type)j->decrement();    }               if(key_compare(key(j), KeyOfValue()(v)))        return pair<link_type,bool>(insert(x,y,v),true);    //Value v was exist    return pair<link_type,bool> (j,false);  }template t_classestypename rb_tree classes::link_type rb_tree classes ::insert_equal(const value_type& v){    link_type y = header;    link_type x = root();    while(x != 0){        y = x;          x = key_compare(KeyOfValue()(v),key(x)) ? left(x) : right(x);    }    return insert(x,y,v);}template t_classestypename rb_tree classes::link_typerb_tree classes ::find(const key_type& k){    link_type cur = root();    if(cur == header) return NULL;    while(cur != NULL){        bool comp1 = key_compare(k,key(cur));        bool comp2 = key_compare(key(cur),k);        if(comp1 == comp2) break;        if(comp1) cur = left(cur);        else cur = right(cur);    }    return cur;}inline voidrb_tree_rotate_left(rb_tree_node_base* x,rb_tree_node_base*& root){    rb_tree_node_base*y = x->right;    x->right = y->left;    if(y->left != NULL)        y->left->parent = x;    y->parent = x->parent;    if(x == root)        root = y;    else if(x == x->parent->left)        x->parent->left = y;    else        x->parent->right = y;    x->left = x;    x->parent = y;}inline voidrb_tree_rotate_right(rb_tree_node_base* x,rb_tree_node_base*& root){    rb_tree_node_base*y = x->left;    x->left = y->right;    if(y->right != NULL)        y->right->parent = x;    y->parent = x->parent;    if(x == root)        root = y;    else if(x == x->parent->right)        x->parent->right = y;    else        x->parent->left = y;    y->right = x;    x->parent = y;}void rb_tree_rebalance(rb_tree_node_base* x,rb_tree_node_base*& root){    x->color = rb_tree_red;     while(x != root && x->parent->color == rb_tree_red){        if(x->parent == x->parent->parent->left){ //parent is the left child of grandpar            rb_tree_node_base* y = x->parent->parent->right;            if(y && y->color == rb_tree_red){ //uncle is red                x->parent->color = rb_tree_black;                y->color = rb_tree_black;                x->parent->parent->color = rb_tree_red;                x = x->parent->parent;            }else{ //no uncle or uncle is black                if(x == x->parent->right){ //x is the right child                    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{ ///parent is the right child of grandpar            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);            }        }    }}void rb_tree_rebalance_for_del(rb_tree_node_base *x,rb_tree_node_base* x_parent,rb_tree_node_base* root){    while (x != root && (x == 0 || x->color == rb_tree_black))         if (x == x_parent->left) {              rb_tree_node_base* w = x_parent->right;              //case 1             if (w->color == rb_tree_red) {                  w->color = rb_tree_black;                  x_parent->color = rb_tree_red;                  rb_tree_rotate_left(x_parent, root);                  w = x_parent->right;              }              //case 2            if ((w->left == NULL || w->left->color == rb_tree_black) &&                  (w->right == NULL || w->right->color == rb_tree_black)) {                  w->color = rb_tree_red;                  x = x_parent;                  x_parent = x_parent->parent;              }               else  {                //case 3                if (w->right == 0 || w->right->color == rb_tree_black) {                      if (w->left) w->left->color = rb_tree_black;                      w->color = rb_tree_red;                      rb_tree_rotate_right(w, root);                      w = x_parent->right;                  }                //case 4                w->color = x_parent->color;                x_parent->color = rb_tree_black;                  if (w->right) w->right->color = rb_tree_black;                  rb_tree_rotate_left(x_parent, root);                  break;              }          }else{            rb_tree_node_base* w = x_parent->left;            //case 1             if (w->color == rb_tree_red) {                w->color = rb_tree_black;                x_parent->color = rb_tree_red;                rb_tree_rotate_right(x_parent, root);                w = x_parent->left;            }            //case 2            if ((w->left == NULL || w->left->color == rb_tree_black) &&                      (w->right == NULL || w->right->color == rb_tree_black)) {                  w->color = rb_tree_red;                       x = x_parent;                x_parent = x_parent->parent;            }            else  {              //case 3                if (w->left == NULL || w->left->color == rb_tree_black) {                    if (w->right) w->right->color = rb_tree_black;                    w->color = rb_tree_red;                    rb_tree_rotate_left(w, root);                    w = x_parent->left;                }                //case 4                w->color = x_parent->color;                x_parent->color = rb_tree_black;                if (w->left) w->left->color = rb_tree_black;                rb_tree_rotate_right(x_parent, root);                break;            }        }         if (x != NULL) x->color = rb_tree_black;  }template t_classesvoid rb_tree classes ::del(link_type& x){    if(x == NULL) return;    if(node_count == 1){        put_node(x);        root() = NULL;        leftmost() = header;        rightmost() = header;        node_count = 0;        return;    }    link_type y = x;    if(right(y) != NULL){        y = right(y);        while(left(y) != NULL) y = left(y);    }    swap(value(x),value(y));    x = y;    if(left(x) != NULL)        y = left(x);    else if(right(x))        y = right(x);    link_type par = parent(x);    if(x != y){        if(x == root()){            root() = y;            leftmost() = y;            rightmost() = y;            parent(y) = parent(x);        }else{            if(x == par->left){                left(par) = y;                parent(y) = par;            }else{                right(par) = y;                parent(y) = par;            }        }    }else{        if(left(par) == x) left(par) = NULL;        else right(par) = NULL;        y = NULL;    }    if(color(x) == rb_tree_black)        rb_tree_rebalance_for_del(y,par,root());    --node_count;    put_node(x);}#endif

简单的测试文件:

#include "rb_tree.h"struct KeyOfValue{    int operator()(int v){return v;}};struct Compare{    bool operator()(const int &a,const int &b){return a<b;}};typedef int Key;typedef int Value;int main(){    rb_tree<Key,Value,KeyOfValue,Compare> tree;    tree.insert_unique(3);    cout<<tree.size()<<endl;    tree.insert_unique(4);    cout<<tree.size()<<endl;    tree.insert_unique(3);    cout<<tree.size()<<endl;    tree.insert_unique(5);    cout<<tree.size()<<endl;    tree.insert_unique(1);    cout<<tree.size()<<endl;    tree.insert_equal(3);    cout<<tree.size()<<endl;    rb_tree<Key,Value,KeyOfValue,Compare>::link_type a = tree.find(3);    rb_tree<Key,Value,KeyOfValue,Compare>::link_type b = tree.find(2);    cout<<b<<int(NULL)<<endl;    tree.del(a);    tree.del(b);    cout<<tree.size()<<endl;return 0;}