红黑树分析与实现之二

下面是红黑树的完整实现，包括插入删除查询等操作的实现，根据需要可以进一步扩展。

由于当初写代码时没有写注释，就直接copy代码过来了，也没加注释。(若真是想钻研代码，结合我的另一篇文章红黑树分析与实现之一还是较容易读懂代码的。)

 

#ifndef __red_black_tree__
#define __red_black_tree__

#define NULL 0

typedef unsigned char color_type;

#define RED 0
#define BLACK 1

template<typename value_type>
struct __rb_tree_node{
 color_type color;
 struct __rb_tree_node *parent;
 struct __rb_tree_node *left;
 struct __rb_tree_node *right;
 value_type value;

 __rb_tree_node():color(RED),parent(NULL),left(NULL),right(NULL){}
 __rb_tree_node(const value_type &v):color(RED),parent(NULL),left(NULL),right(NULL),value(v){}

 struct __rb_tree_node *deep_copy() const
 {
  struct __rb_tree_node *copy = new __rb_tree_node(*this);
  copy->parent = NULL;
  if (this->left)
  {
   copy->left = this->left->deep_copy();
   if (copy->left)
    copy->left->parent = copy;
  }
  if (this->right)
  {
   copy->right = this->right->deep_copy();
   if (copy->right)
    copy->right->parent = copy;
  }
  return copy;
 }

 ~__rb_tree_node(){
  if (left)
   delete left;
  left = NULL;
  if (right)
   delete right;
  right = NULL;
  parent = NULL;
 }
};

template<typename value_type>
__rb_tree_node<value_type> *__rb_tree_node_left_rotate(struct __rb_tree_node<value_type> *node)
{
 if (node==NULL||node->right==NULL)
  return NULL;

 struct __rb_tree_node<value_type> *right = node->right;

 node->right = right->left;
 if (node->right)
  node->right->parent = node;

 right->left = node;

 right->parent = node->parent;

 if (node->parent)
  if (node==node->parent->left)
   node->parent->left = right;
  else
   node->parent->right = right;
 
 node->parent = right;

 return right;
}

template<typename value_type>
__rb_tree_node<value_type> *__rb_tree_node_right_rotate(struct __rb_tree_node<value_type> *node)
{
 if (node==NULL||node->left==NULL)
  return NULL;

 struct __rb_tree_node<value_type> *left = node->left;

 node->left = left->right;
 if (node->left)
  node->left->parent = node;

 left->right = node;

 left->parent = node->parent;

 if (node->parent)
  if (node==node->parent->left)
   node->parent->left = left;
  else
   node->parent->right = left;
 
 node->parent = left;

 return left;
}

template<typename value_type>
__rb_tree_node<value_type> *__rb_tree_node_first(struct __rb_tree_node<value_type> *node)
{
 if (node==NULL)
  return NULL;

 struct __rb_tree_node<value_type> *first = node;

 while (first->left)
  first = first->left;

 return first;
}

template<typename value_type>
__rb_tree_node<value_type> *__rb_tree_node_last(struct __rb_tree_node<value_type> *node)
{
 if (node==NULL)
  return NULL;

 struct __rb_tree_node<value_type> *last = node;

 while (last->right)
  last = last->right;

 return last;
}

template<typename value_type>
__rb_tree_node<value_type> *__rb_tree_node_next(struct __rb_tree_node<value_type> *node)
{
 if (node==NULL)
  return NULL;

 struct __rb_tree_node<value_type> *next = node->right;

 if (next)
 {
  while (next->left)
   next = next->left;
 }
 else
 {
  struct __rb_tree_node<value_type> *parent = node->parent;
  struct __rb_tree_node<value_type> *child = node;
  while (parent&&parent->right==child)
  {
   child = parent;
   parent = parent->parent;
  }

  next = parent;
 }

 return next;
}

template<typename value_type>
__rb_tree_node<value_type> *__rb_tree_node_previous(struct __rb_tree_node<value_type> *node)
{
 if (node==NULL)
  return NULL;

 struct __rb_tree_node<value_type> *pre = node->left;

 if (pre)
 {
  while (pre->right)
   pre = pre->right;
 }
 else
 {
  struct __rb_tree_node<value_type> *parent = node->parent;
  struct __rb_tree_node<value_type> *child = node;
  while (parent&&parent->left==child)
  {
   child = parent;
   parent = parent->parent;
  }

  pre = parent;
 }

 return pre;
}

template<typename value_type>
int __rb_tree_node_height(struct __rb_tree_node<value_type> *node)
{
 if (node==NULL)
  return 0;

 int lh = ::__rb_tree_node_height(node->left);
 int rh = ::__rb_tree_node_height(node->right);
 return 1 + (lh>rh?lh:rh);
}

template<typename value_type>
class RB_Tree{
public:
 typedef struct __rb_tree_node<value_type> NODE;
 typedef struct __rb_tree_node<value_type> *PNODE;
 typedef value_type VALUE_TYPE;

private:
 PNODE root;
 int _size;
public:
 RB_Tree():root(NULL),_size(0){}
 RB_Tree(const RB_Tree& other){
  if (other.root)
   this->root = other.root->deep_copy();
  else
   this->root = NULL;
  this->_size = other._size;
 }

 RB_Tree &operator=(const RB_Tree& other){
  if (this==&other)
   return *this;

  if (this->root)
   delete this->root;

  if (other.root)
   this->root = other.root->deep_copy();
  else
   this->root = NULL;

  this->_size = other._size;
  return *this;
 }

 ~RB_Tree(){
  if (root)
   delete root;
  root = NULL;
 }

 PNODE first()
 {
  return ::__rb_tree_node_first(this->root);
 }

 PNODE last()
 {
  return ::__rb_tree_node_last(this->root);
 }

 PNODE next(PNODE p)
 {
  return ::__rb_tree_node_next(p);
 }
 
 PNODE previous(PNODE p)
 {
  return ::__rb_tree_node_previous(p);
 }

 PNODE search(const VALUE_TYPE &value)
 {
  PNODE p = root;
  while (p)
  {
   if (value<p->value)
    p = p->left;
   else if (p->value<value)
    p = p->right;
   else
   {
    PNODE tmp = this->previous(p);
    while (tmp&&!(tmp->value<value))
    {
     p = tmp;
     tmp = this->previous(p);
    }
    break;
   }
  }
  return p;
 }

 void add(const VALUE_TYPE &value)
 {
  PNODE p = root;
  while (p)
  {
   if (p->value<value&&p->right)
    p = p->right;
   else if (value<p->value&&p->left)
    p = p->left;
   else
    break;
  }
  if (p==NULL)
  {
   root = new NODE(value);
   root->color = BLACK;
  }
  else
  {
   PNODE an = new NODE(value);
   if (value<p->value)
   {
    p->left = an;
    p->left->parent = p;
   }
   else if (p->value<value)
   {
    p->right = an;
    p->right->parent = p;
   }
   else
   {
    PNODE next = this->next(p);
    while (next)
    {
     p = next;
     if (value<next->value)
      break;
     else
      next = this->next(p);
    }
    if (value<p->value)
    {
     an->value = p->value;
     p->value = value;
     PNODE f = ::__rb_tree_node_first(p->right);
     if (f)
     {
      p = f;
      p->left = an;
      p->left->parent = p;
     }
     else
     {
      p->right = an;
      p->right->parent = p;
     }
    }
    else
    {
     p->right = an;
     p->right->parent = p;
    }
   }
   if (p->color==RED)
   {
    add_adjust(an);
   }
  }
  ++this->_size;
 }

 bool check(int &height)
 {
  return this->check(this->root,height);
 }

 void remove(const VALUE_TYPE & value)
 {
  PNODE itr = this->search(value);
  while (itr)
  {
   PNODE tmp = this->next(itr);
   this->remove_internal(itr);
   itr = tmp&&tmp->value==value?tmp:NULL;
  }
 }

 int size(){return this->_size;}

 int height(){return ::__rb_tree_node_height(this->root);}

private:
 void remove_internal(PNODE p)
 {
  if (!p)
   return;

  if (p->left&&p->right)
  {
   PNODE pre = ::__rb_tree_node_previous(p);
   p->value = pre->value;
   p = pre;
  }

  PNODE parent = p->parent;
  PNODE child = NULL;
  if (p->left)
   child = p->left;
  else
   child = p->right;

  if (parent)
  {
   if (p==parent->left)
    parent->left = child;
   else
    parent->right = child;
  }
  else
  {
   this->root = child;
  }

  if (child)
  {
   child->parent = parent;
  }

  if (p->color==BLACK)
  {
   this->remove_adjust(parent,child);
  }

  p->parent = NULL;
  p->left = NULL;
  p->right = NULL;
  delete p;
  p = NULL;

  --this->_size;
 }

 color_type color_of(PNODE p)
 {
  if (!p)
   return BLACK;
  return p->color;
 }

 void remove_adjust(PNODE parent,PNODE child)
 {
  if (parent==NULL)//root
  {
   this->root = child;
   if (this->root)
    this->root->color = BLACK;
   return;
  }

  if (child&&child->color==RED)
  {
   child->color = BLACK;
   return;
  }

  if (parent->left==child)
  {
   if (parent->right->color==RED)
   {
    parent->color = RED;
    parent->right->color = BLACK;
    PNODE np = ::__rb_tree_node_left_rotate(parent);
    if (parent==this->root)
     this->root = np;
    parent = np->left;
    this->remove_adjust(parent,child);
   }
   else
   {
    if (this->color_of(parent->right->right)==RED)
    {
     parent->right->right->color = BLACK;
     parent->right->color = parent->color;
     parent->color = BLACK;
     PNODE np = ::__rb_tree_node_left_rotate(parent);
     if (parent==this->root)
      this->root = np;
    }
    else if (this->color_of(parent->right->left)==BLACK)
    {
     parent->right->color = RED;
     this->remove_adjust(parent->parent,parent);
    }
    else
    {
     parent->right->color = RED;
     parent->right->left->color = BLACK;
     PNODE np = ::__rb_tree_node_right_rotate(parent->right);
     if (parent==this->root)
      this->root = np;
     this->remove_adjust(parent,child);
    }
   }
  }
  else
  {
   if (parent->left->color==RED)
   {
    parent->color = RED;
    parent->left->color = BLACK;
    PNODE np = ::__rb_tree_node_right_rotate(parent);
    if (parent==this->root)
     this->root = np;
    parent = np->right;
    this->remove_adjust(parent,child);
   }
   else
   {
    if (this->color_of(parent->left->left)==RED)
    {
     parent->left->left->color = BLACK;
     parent->left->color = parent->color;
     parent->color = BLACK;
     PNODE np = ::__rb_tree_node_right_rotate(parent);
     if (parent==this->root)
      this->root = np;
    }
    else if (this->color_of(parent->left->right)==BLACK)
    {
     parent->left->color = RED;
     this->remove_adjust(parent->parent,parent);
    }
    else
    {
     parent->left->color = RED;
     parent->left->right->color = BLACK;
     PNODE np = ::__rb_tree_node_left_rotate(parent->left);
     if (parent==this->root)
      this->root = np;
     this->remove_adjust(parent,child);
    }
   }
  }
 }

 void add_adjust(PNODE node)
 {
  if (node==NULL||node->parent==NULL||node->color!=RED||node->parent->color!=RED)
   return;

  PNODE gp = node->parent->parent;

  if (node->parent==gp->left)
  {
   if (gp->right&&gp->right->color==RED)
   {
    gp->color = RED;
    gp->left->color = BLACK;
    gp->right->color = BLACK;
    if (gp->parent&&gp->parent->color==RED)
     add_adjust(gp);
    else if (gp->parent==NULL)//gp是root
     gp->color = BLACK;
   }
   else
   {
    if (node==node->parent->right)
     ::__rb_tree_node_left_rotate(node->parent);
    gp->color = RED;
    gp->left->color = BLACK;
    PNODE nr = ::__rb_tree_node_right_rotate(gp);
    if (gp==root)
     root = nr;
   }
  }
  else
  {
   if (gp->left&&gp->left->color==RED)
   {
    gp->color = RED;
    gp->left->color = BLACK;
    gp->right->color = BLACK;
    if (gp->parent&&gp->parent->color==RED)
     add_adjust(gp);
    else if (gp->parent==NULL)//gp是root
     gp->color = BLACK;
   }
   else
   {
    if (node==node->parent->left)
     ::__rb_tree_node_right_rotate(node->parent);
    gp->color = RED;
    gp->right->color = BLACK;
    PNODE nr = ::__rb_tree_node_left_rotate(gp);
    if (gp==root)
     root = nr;
   }
  }
 }

 bool check(PNODE node,int &height)
 {
  height = 0;
  if (!node)
   return true;

  bool ok = true;
  if (node->color==BLACK)
  {
   height = 1;
  }
  else
  {
   if (node->left&&node->left->color==RED)
    ok = false;
   else if (node->right&&node->right->color==RED)
    ok = false;
  }

  if (ok)
  {
   int lh = 0;
   int rh = 0;
   if (!(check(node->left,lh)&&check(node->right,rh)&&lh==rh))
   {
    ok = false;
   }
   else
   {
    height += lh;
   }
  }

  return ok;
 }
};
#endif