红黑树源码实现

测试

原理不说了，参考下面的链接：

http://hi.baidu.com/20065562/blog/item/93b2d17fd6f391320dd7da44.html

 

目前实现了添加和list操作。（删除有点类似，写代码太复杂，就写这么多了）

 

最后用dot命令把红黑树给画了出来,效果还行。（dot命令能够是graphviz开源软件提供的）

 

展示一下效果：

10个元素，正序分别插入1，2，3，4。。。之后的红黑树图。

 

 

10个元素，逆序分别插入10，9，8，7。。。之后的红黑树图。

 

10个元素，乱序分别插入。。。。之后的红黑树图。

 

dot命令：

./dot rb_tree.dot -Tpng -o rb_tree.png

rb_tree.dot 文件如下：

 

digraph g {
  {
    node
    [
     style=filled,
     fontcolor="white",
     color="red"
    ] 1 3
    node
    [
     fontcolor="white",
     style=filled,
     color="black"
    ]2 4 5 6 7 8 9 10
  }
  7 -> 5
  5 -> 3
  3 -> 2
  2 -> 1
  3 -> 4
  5 -> 6
  7 -> 9
  9 -> 8
  9 -> 10
 }


 

rb_tree.cc 文件如下：

 

#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#include <vector>
#include <iostream>

using namespace std;


vector<int> vec_red_nums;
vector<int> vec_black_nums;

typedef enum color_t
{
  RED,
  BLACK
} color_t;

typedef struct rb_node_t
{
  struct rb_node_t *parent;
  struct rb_node_t *left_child;
  struct rb_node_t *right_child;
  color_t color;
  int key;
  int data;
} rb_node_t;

typedef struct rb_tree_t
{
  struct rb_node_t *root;
  int count;
} rb_tree_t;

void left_rotate(rb_node_t *node)
{
   cout << "l left rotate node " << node->key << endl;
   rb_node_t *parent = node->parent;
   rb_node_t *rchild = node->right_child;

   if (parent != NULL)
   parent->left_child = rchild;

   rchild->parent = parent;

   node->right_child = rchild->left_child;
   node->parent = rchild;

   rchild->left_child = node;
   return;
}

void r_left_rotate(rb_node_t *node)
{
   cout << "r left rotate node " << node->key << endl;
   rb_node_t *parent = node->parent;
   rb_node_t *rchild = node->right_child;

   if (parent != NULL)
   parent->right_child = rchild;

   rchild->parent = parent;

   node->right_child = rchild->left_child;
   node->parent = rchild;

   rchild->left_child = node;
   return;
}


void right_rotate(rb_node_t *node)
{
  cout << "l right rotate node " << node->key << endl;
  rb_node_t *parent = node->parent;
  rb_node_t *lchild = node->left_child;

  if (parent != NULL)
  parent->left_child = lchild;

  lchild->parent = parent;

  node->left_child = lchild->right_child;
  node->parent = lchild;

  lchild->right_child = node;
  return;
}

void r_right_rotate(rb_node_t *node)
{
  cout << "r right rotate node " << node->key << endl;
  rb_node_t *parent = node->parent;
  rb_node_t *lchild = node->left_child;

  if (parent != NULL)
  parent->right_child = lchild;

  lchild->parent = parent;

  node->left_child = lchild->right_child;
  node->parent = lchild;

  lchild->right_child = node;
  return;
}


rb_node_t *find_node(int key, rb_tree_t *tree)
{
  rb_node_t *p = tree->root;
  while (NULL != p)
  {
    if (key < p->key)
    {
      p = p->left_child;
    }
    else if (key > p->key)
    {
      p = p->right_child;
    }
    else
    {
      return p;
    }
  }
  return NULL;
}

rb_node_t *add_node(int key, int data, rb_tree_t *tree)
{
  rb_node_t *node = (rb_node_t*)malloc(sizeof(rb_node_t));
  node->left_child = NULL;
  node->right_child = NULL;
  node->color = RED;
  node->key = key;
  node->data = data;

  rb_node_t *p = tree->root;
  rb_node_t *p_pre = tree->root;
  while (NULL != p)
  {
    p_pre = p;
    if (key < p->key)
    {
      p = p->left_child;
    }
    else
    {
      p = p->right_child;
    }
  }

  if (key < p_pre->key)
  {
    p_pre->left_child = node;
  }
  else
  {
    p_pre->right_child = node;
  }

  node->parent = p_pre;
  //cout << "node " << key << " parent = " << node->parent << endl;

  tree->count++;

  return node;
}

void insert_node(int key, int data, rb_tree_t *tree)
{
  //empty tree

  if (tree->root == NULL)
  {
    rb_node_t *node = (rb_node_t*)malloc(sizeof(rb_node_t));
    if (NULL == node)
    {
      printf("malloc failed\n");
      return;
    }
    node->parent = NULL;
    node->left_child = NULL;
    node->right_child = NULL;
    node->color = BLACK;
    node->key = key;
    node->data = data;

    tree->root = node;
    tree->count = 1;
    return;
  }

  //node already exsit

  rb_node_t *p = find_node(key, tree);
  if (p != NULL)
  {
    p->data = data;
    return;
  }

  //add node

  rb_node_t *node = add_node(key, data, tree);

  if (tree->count == 1)return;

  //do balance

  while (true)
  {
    if (node == tree->root)
    {
      node->color = BLACK;
      break;
    }

    if (node->parent->color == BLACK)
    {
      break;
    }

    if (node->parent == node->parent->parent->left_child) //left sub tree

    {
      if (node->parent->color == RED) //need balance

      {
        if (node->parent->parent->right_child != NULL && node->parent->parent->right_child->color == RED) //red uncle

        {
          node->parent->color = BLACK;
          node->parent->parent->right_child->color = BLACK;
          node->parent->parent->color = RED;
          node = node->parent->parent;
        }
        else //black uncle

        {
          if (node == node->parent->left_child) //node is left child

          {
            node->parent->color = BLACK;
            node->parent->parent->color = RED;
            if (node->parent->parent == tree->root)
            {
              tree->root = node->parent;
            }
            right_rotate(node->parent->parent);
          }
          else //node is right child

          {
            node->color = BLACK;
            node->parent->parent->color = RED;
            if (node->parent->parent == tree->root)
            {
              tree->root = node;
            }
            left_rotate(node->parent);
            right_rotate(node->parent);
          }
        }
      }
      else
      {
        break;
      }
    }
    else //right sub tree

    {
      if (node->parent->color == RED) //need balance

      {
        if (node->parent->parent->left_child != NULL && node->parent->parent->left_child->color == RED) //red uncle

        {
          node->parent->color = BLACK;
          node->parent->parent->left_child->color = BLACK;
          node->parent->parent->color = RED;
          node = node->parent->parent;
        }
        else //black uncle

        {
          //四种情况，重新区分

          if (node == node->parent->right_child) //node is right child

          {
            node->parent->color = BLACK;
            node->parent->parent->color = RED;
            if (node->parent->parent == tree->root)
            {
              tree->root = node->parent;
            }
            r_left_rotate(node->parent->parent);
          }
          else //node is left child

          {
            node->color = BLACK;
            node->parent->parent->color = RED;
            if (node->parent->parent == tree->root)
            {
              tree->root = node;
            }
            r_right_rotate(node->parent);
            r_left_rotate(node->parent);
          }
        }
      }
      else
      {
        break;
      }
    }
  }
  return;
}

void list_node(rb_node_t *node)
{
  if (node->left_child != NULL)
  {
    printf("%d -> %d\n", node->key, node->left_child->key);
    list_node(node->left_child);
  }

  if (node->color == BLACK)
  {
    vec_black_nums.push_back(node->key);
  }
  else
  {
    vec_red_nums.push_back(node->key);
  }

  if (node->right_child != NULL)
  {
    printf("%d -> %d\n", node->key, node->right_child->key);
    list_node(node->right_child);
  }
  return;
}

void list_tree(rb_tree_t *tree)
{
  vec_red_nums.clear();
  vec_black_nums.clear();
  printf("tree nodes count = %d\n", tree->count);
  list_node(tree->root);

  cout << "red = ";
  for (int i=0; i<vec_red_nums.size(); i++)
  {
    cout << vec_red_nums[i] << " ";
  }
  cout << endl;

  cout << "black = ";
  for (int i=0; i<vec_black_nums.size(); i++)
  {
    cout << vec_black_nums[i] << " ";
  }
  cout << endl;
  return;
}

int main(int argc, char **argv)
{

  int num = atoi(argv[1]);

  int key_arr[10] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
  int data_arr[10] = {10, 20, 30, 40, 50, 60, 70, 80, 90, 100};
  //int indexs[10] = {7, 1, 2, 5, 9, 0, 4, 3, 8, 6};

  //int indexs[10] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};

  int indexs[10] = {9, 8, 7, 6, 5, 4, 3, 2, 1, 0};

  rb_tree_t tree;
  tree.root = NULL;

  int i;
  for (i=0; i<num; i++)
  {
    cout << "insert node " << key_arr[indexs[i]] << endl;
    insert_node(key_arr[indexs[i]], data_arr[indexs[i]], &tree);
  }

  list_tree(&tree);

  return 0;
}