红黑树的代码实现

红黑树实际上是一种高效的平衡二叉树，它的平衡性不如avl树严格，但是统计性能更高，在对大数据的信息进行排序查找操作时，如果对效率的要求不是特别的高，红黑树是一个恰当的选择，相对于散列表，红黑树的“代码复杂程度”更低，因为散列表要牵涉到散列函数，一个优越的散列函数并不是那么容易构造。关于红黑树与AVL树的对比，请参阅：http://wangdei.iteye.com/blog/236157

下面是红黑树实现的代码，代码经过测试。

/*********************************************************************************
*rb_tree.h
*Version:  
*Date:  July 2013
*Description:  insert，search，free function of rbTree
*History:  none
**********************************************************************************/ 
#ifndef  RBTREE
#define  RBTREE
#include <stdio.h>
#include <iostream>

//
//data type define
//
extern struct Triangle;


struct ElemType
{
    ElemType()
    {
    }
    ~ElemType()
    {
    }
    float* coor;

    bool operator >( const ElemType& e1 ) const
    {
        if ( coor[0] > e1.coor[0] )
        {
            return true;
        }
        else if ( coor[0] == e1.coor[0] && coor[1] > e1.coor[1] )
        {
            return true;
        }
        else if ( coor[0] == e1.coor[0] && coor[1] == e1.coor[1] && coor[2] > e1.coor[2] )
        {
            return true;
        }
        else
        {
            return false;
        }
    }

    bool operator <( const ElemType& e1 ) const
    {
        if ( coor[0] < e1.coor[0] )
        {
            return true;
        }
        else if ( coor[0] == e1.coor[0] && coor[1] < e1.coor[1] )
        {
            return true;
        }
        else if ( coor[0] == e1.coor[0] && coor[1] == e1.coor[1] && coor[2] < e1.coor[2] )
        {
            return true;
        }
        else
        {
            return false;
        }
    }

    bool operator ==( const ElemType& e1 ) const
    {
        if ( coor[0] == e1.coor[0] && coor[1] == e1.coor[1] && coor[2] == e1.coor[2] )
        {
            return true;
        }
        else
        {
            return false;
        }
    }

    ElemType& operator=( const ElemType& e1 )
    {
        coor = e1.coor;

        return *this;
    }
}; 


//红黑树节点类
typedef struct rb_node
{
    rb_node( ElemType* data )
    {
        this->data = data;
    }

    rb_node()
    {
        this->data = 0;
    }

    ~rb_node()
    {
    }

    ElemType* data; 
    char color;//R,B--red,black
    struct rb_node* pa;//parent
    struct rb_node* left;    
    struct rb_node* right;
}rb_node, rb_tree;



class RB_Tree
{
  public:
    RB_Tree()
    {
        rb_root = new rb_tree( 0 );

        rb_root->pa = rb_root;
        rb_root->left = rb_root;
        rb_root->right = rb_root;
        rb_root->color = 'B';

        size = 0;
    }

    ~RB_Tree()
    {
        delete rb_root;

        rb_root = 0;
    }

    int rb_tree_insert( rb_node* node );//插入节点

    void rb_get_node( rb_node* node );//设置空节点

    int rb_tree_fixup( rb_node* node );//修正操作（左旋，右旋）

    rb_node* rb_tree_init( rb_node* a );//初始化红黑树

    int rbtree_output( char* filePath );//for test

    bool rbtree_search( ElemType& key ); //search element

    int free_rbtree();//清空

    int shallow_free_rbtree();//浅清空，指针置为零，不删除内容

    rb_node* tree_successor( rb_node* node );//node的后继

    rb_node* tree_first();//返回第一个节点

    int rb_get_array( rb_node* node_array, int& count );//前序遍历返回所以节点

    int size;

  private:
    rb_tree* rb_root;



    int rb_tree_insert( rb_tree* tree, rb_node* node );

    int free_rbtree( rb_tree* );

    int rbtree_output( char*& mainchar, rb_tree* tree, int level );

    bool rbtree_search( rb_tree* node, ElemType& key );

    int rbtree_left_rotate( rb_tree* tree, rb_node* node );

    int rbtree_right_rotate( rb_tree* tree, rb_node* node );

    int rb_tree_fixup( rb_tree* tree, rb_node* node );

    int os_key_rank( rb_tree* tree, ElemType& key );
    int os_rank( rb_tree* tree, rb_node* node );
    int right_tree_num( rb_node* node );
    int left_tree_num( rb_node* node );

    int shallow_free_rbtree( rb_tree* );

    int rb_get_array( rb_node* node_array, int& count, rb_node node );
};

#endif

#include "rb_tree.h"


int RB_Tree::free_rbtree()
{
  rb_tree* tree = rb_root;

  if ( tree->left != rb_root )
  {
    free_rbtree( tree->left );
  }

  if ( tree->right != rb_root )
  {
    free_rbtree( tree->right );
  }

  delete tree;
  tree = 0;
  return 0;
}

rb_node* RB_Tree::rb_tree_init( rb_node* root )
{
  root->color = 'B';
  root->left = rb_root;//nil
  root->right = rb_root;
  root->pa = rb_root; 
  rb_root->left = rb_root->right = root;

  size = 1;

  return rb_root->left;
}

int RB_Tree::free_rbtree( rb_tree* tree )
{
  if ( tree->left != rb_root )
  {
    free_rbtree( tree->left );
  }

  if ( tree->right != rb_root )
  {
    free_rbtree( tree->right );
  }

  delete tree;
  tree = 0;
  return 0;
}

int RB_Tree::left_tree_num( rb_node* node )
{
  int left, right, leftnum;
  if ( node == rb_root )
  {
    return 0;
  }
  else if ( node->left != rb_root )
  {
    left = left_tree_num( node->left );
    right = right_tree_num( node->left );
    leftnum = left + right + 1;
  }
  else
  {
    return 0;
  }
  return leftnum;
}

int RB_Tree::right_tree_num( rb_node* node )
{
  int left, right, rightnum;
  if ( node == rb_root )
  {
    return 0;
  }
  else if ( node->right != rb_root )
  {
    left = left_tree_num( node->right );
    right = right_tree_num( node->right );
    rightnum = left + right + 1;
  }
  else
  {
    return 0;
  }
  return rightnum;
}


int RB_Tree::os_rank( rb_tree* tree, rb_node* node )
{
  rb_node* y;
  int rank;
  rank = left_tree_num( node ) + 1;
  y = node;
  while ( y != rb_root->left )
  {
    if ( y == y->pa->right )
    {
      rank = rank + left_tree_num( y->pa ) + 1;
    }
    y = y->pa;
  }
  return rank;
}

int RB_Tree::os_key_rank( rb_tree* tree, ElemType& key )
{
  int num;
  rb_node* node;
  node = tree;
  while ( !( *( node ->data ) == key ) )
  {
    if ( *( node->data ) > key )
    {
      node = node ->left;
    }
    else
    {
      node = node ->right;
    }
  }
  if ( node == rb_root )
  {
    return -1;
  }
  num = os_rank( tree, node );
  return num;
}





//LDR中序遍历输出
int RB_Tree::rbtree_output( char*& mainchar, rb_tree* tree, int level )
{
  int i = 0,lev = level;

  char* tempchar = new char[1024];

  //  fprintf(result,"(%d%c, \n",tree->data,tree->color);
  sprintf( tempchar, "(%f %f %f %c, \n", tree->data->coor[0], tree->data->coor[1], tree->data->coor[2], tree->color );

  strcat( mainchar, tempchar );

  if ( tree->left != rb_root )
  {
    //rb_root = nil
    for ( i = 0; i <= level; i++ )
    {
      //fprintf(result,"\t");
      sprintf( tempchar, "\t" );
      strcat( mainchar, tempchar );
    }
    rbtree_output( mainchar, tree->left, ++level );
  }
  else
  {
    for ( i = 0; i <= level; i++ )
    {
      sprintf( tempchar, "\t" );
      strcat( mainchar, tempchar );
    }
    sprintf( tempchar, "NIL,\n" );
    strcat( mainchar, tempchar );
    //    fprintf(result,"NIL,\n");
  }
  //  system("pause");
  if ( tree->right != rb_root )
  {
    for ( i = 0; i <= lev; i++ )
    {
      sprintf( tempchar, "\t" );
      strcat( mainchar, tempchar );
    }
    rbtree_output( mainchar, tree->right, ++lev );
  }
  else
  {
    for ( i = 0; i <= lev; i++ )
    {
      sprintf( tempchar, "\t" );
      strcat( mainchar, tempchar );
    }
    sprintf( tempchar, "NIL)\n" );
    strcat( mainchar, tempchar );
  }
  delete[] tempchar;
  return 0;
}

//LDR中序遍历输出
int RB_Tree::rbtree_output( char* filePath )
{
  FILE* fp;

  if ( ( fp = fopen( filePath, "w+" ) ) == NULL )
  {
    printf( "can not open this file\n" );

    return false;
  }

  char* mainchar = new char[1024000];

  char* tempchar = new char[1024];


  memset( mainchar, 0, sizeof( char ) * 1024000 );
  memset( tempchar, 0, sizeof( char ) * 1024 );

  rb_tree* tree = rb_root->left;
  int level = 0;


  //right & left
  int i = 0,lev = level;



  //  fprintf(result,"(%d%c, \n",tree->data,tree->color);
  sprintf( tempchar, "(%f %f %f %c, \n", tree->data->coor[0], tree->data->coor[1], tree->data->coor[2], tree->color );
  strcat( mainchar, tempchar );
  if ( tree->left != rb_root )
  {
    //rb_root = nil
    for ( i = 0; i <= level; i++ )
    {
      //fprintf(result,"\t");
      sprintf( tempchar, "\t" );
      strcat( mainchar, tempchar );
    }
    rbtree_output( mainchar, tree->left, ++level );
  }
  else
  {
    for ( i = 0; i <= level; i++ )
    {
      sprintf( tempchar, "\t" );
      strcat( mainchar, tempchar );
    }
    sprintf( tempchar, "NIL,\n" );
    strcat( mainchar, tempchar );

    //    fprintf(result,"NIL,\n");
  }
  //  system("pause");
  if ( tree->right != rb_root )
  {
    for ( i = 0; i <= lev; i++ )
    {
      sprintf( tempchar, "\t" );
      strcat( mainchar, tempchar );
    }
    rbtree_output( mainchar, tree->right, ++lev );
  }
  else
  {
    for ( i = 0; i <= lev; i++ )
    {
      sprintf( tempchar, "\t" );
      strcat( mainchar, tempchar );
    }
    sprintf( tempchar, "NIL)\n" );
    strcat( mainchar, tempchar );
  }


  fputs( mainchar, fp );

  delete[] tempchar;

  delete[] mainchar;

  fclose( fp );

  return 0;
}






rb_node* RB_Tree::tree_first()//node的后继
{
  rb_node* p;

  p = rb_root->left;

  if ( p == rb_root )
  {
    return NULL;
  }

  return p;
}

rb_node* RB_Tree::tree_successor( rb_node* node )//node的后继
{
  rb_node* p;

  p = node->right;

  if ( p == rb_root )
  {
    return NULL;
  }

  while ( p->left != rb_root )
  {
    p = p->left;
  }

  return p;
}
int RB_Tree::rb_tree_insert( rb_tree* tree, rb_node* node )
{
  if ( node == NULL || tree == NULL )
  {
    return 0;
  }

  rb_node* p = NULL;
  rb_node* t = tree;   

  while ( t != rb_root )
  {
    p = t ;

    if ( *( node->data ) < *( t->data ) )
    {
      t = t->left;
    }
    else if ( *( node->data ) > *( t->data ) )
    {
      t = t->right;
    }
    else if ( *( node->data ) == *( t->data ) )
    {
      return 0;
    }
  }

  if ( *( node->data ) < *( p->data ) )
  {
    p->left = node;
  }
  else
  {
    p->right = node;
  }

  size++;

  node->pa = p;
  node->color = 'R';

  return 1;
}


int RB_Tree::rb_get_array( rb_node* node_array, int& count )
{
  if ( rb_root->left == rb_root )
  {
    return 0;
  }

  count = 0;

  node_array[count] = *( rb_root->left );

  count++;

  if ( rb_root->left->left != rb_root )
  {
    rb_get_array( node_array, count, *( rb_root->left->left ) );
  }

  if ( rb_root->left->right != rb_root )
  {
    rb_get_array( node_array, count, *( rb_root->left->right ) );
  }

  return 1;
}

int RB_Tree::rb_get_array( rb_node* node_array, int& count, rb_node src_node )
{
  node_array[count] = src_node;

  count++;

  if ( src_node.left != rb_root )
  {
    rb_get_array( node_array, count, *( src_node.left ) );
  }

  if ( src_node.right != rb_root )
  {
    rb_get_array( node_array, count, *( src_node.right ) );
  }

  return 1;
}


int RB_Tree::rb_tree_insert( rb_node* node )
{
  rb_get_node( node );
  return rb_tree_insert( rb_root->left, node );
}


int RB_Tree::rb_tree_fixup( rb_node* node )
{
  return rb_tree_fixup( rb_root->left, node );
}


int RB_Tree::rb_tree_fixup( rb_tree* tree, rb_node* node )
{
  while ( node->pa->color == 'R' )
  {
    if ( node->pa == node->pa->pa->left )
    {
      if ( node->pa->pa->right->color == 'R' )
      {
        node = node->pa->pa;
        node->left->color = node->right->color = 'B';
        node->color = 'R';
      }
      else
      {
        if ( node->pa->right == node )
        {
          node = node->pa; 
          rbtree_left_rotate( tree, node );
        }
        node->pa->color = 'B';
        node->pa->pa->color = 'R';
        rbtree_right_rotate( tree, node->pa->pa );
      }
    }
    else
    {
      if ( node->pa->pa->left->color == 'R' )
      {
        node = node->pa->pa;
        node->left->color = node->right->color = 'B'; 
        node->color = 'R';
      }
      else
      {
        if ( node->pa->left == node )
        {
          node = node->pa;

          rbtree_right_rotate( tree, node );
        } 

        node->pa->color = 'B'; 
        node->pa->pa->color = 'R';
        rbtree_left_rotate( tree, node->pa->pa );
      }
    }
  }
  rb_root->left->color = 'B';
  return 0;
}


int RB_Tree::rbtree_left_rotate( rb_tree* tree, rb_node* node )
{
  rb_node* c;

  c = node->right;

  if ( node->right == rb_root )
  {
    return 1;
  }

  if ( node->pa != rb_root )
  {
    if ( node->pa->left == node )
    {
      node->pa->left = c;
    }
    else
    {
      node->pa->right = c;
    }
  }
  else
  {
    rb_root->left = c;
  }
  c->pa = node->pa;

  node->right = c->left;
  if ( c->left != rb_root )
  {
    c->left->pa = node;
  }
  c->left = node;
  node->pa = c; 

  return 0;
}

int RB_Tree::rbtree_right_rotate( rb_tree* tree, rb_node* node )
{
  rb_node* c;   

  c = node->left;

  if ( node->left == rb_root )
  {
    return 1;
  }    

  if ( node->pa != rb_root )
  {
    if ( node->pa->left == node )
    {
      node->pa->left = c;
    }
    else
    {
      node->pa->right = c;
    }
  }
  else
  {
    rb_root->left = c;
  }
  c->pa = node->pa;   
  node->left = c->right;
  if ( c->right != rb_root )
  {
    c->right->pa = node;
  }
  c->right = node;
  node->pa = c;

  return 0;
}

bool RB_Tree::rbtree_search( ElemType& key )
{
  rb_tree* tree = rb_root->left;
  return rbtree_search( tree, key );
}

bool RB_Tree::rbtree_search( rb_tree* tree, ElemType& key )
{
  if ( rb_root->left == rb_root )
  {
    return false;
  }

  rb_node* node;
  node = rb_root->left;//rb_root->left 指向根，所有叶子指向rb_root;

  if ( node != tree )
  {
    printf( "rb_tree error\n" );
    return false;
  }

  while ( node != rb_root )
  {
    if ( *( node->data ) < key )
    {
      node = node->right;
    }
    else if ( *( node->data ) > key )
    {
      node = node->left;
    }
    else if ( *( node->data ) == key )
    {
      break;
    }
  }
  if ( node == rb_root )
  {
    printf( "No key %f %f %f in rbtree!\n", key.coor[0], key.coor[1], key.coor[2] );

    tree = rb_root->left;

    return false;
  }

  tree = node;

  return true;
}

void RB_Tree::rb_get_node( rb_node* node )
{
  node->color = 'W';
  node->pa = rb_root;
  node->left = rb_root;
  node->right = rb_root;
}

int RB_Tree::shallow_free_rbtree()
{
  rb_tree* tree = rb_root;

  if ( tree->left != rb_root )
  {
    shallow_free_rbtree( tree->left );
  }

  if ( tree->right != rb_root )
  {
    shallow_free_rbtree( tree->right );
  }

  tree = NULL;
  return 0;
}

int RB_Tree::shallow_free_rbtree( rb_tree* tree )
{
  if ( tree->left != rb_root )
  {
    shallow_free_rbtree( tree->left );
  }

  if ( tree->right != rb_root )
  {
    shallow_free_rbtree( tree->right );
  }

  tree = NULL;

  return 0;
}