a implementation of red-black tree in cpp

rbtree.cpp

/*
 * =====================================================================================
 *
 *       Filename:  rbtree.cpp
 *
 *    Description:  red-black tree in cpp
 *
 *        Version:  1.0
 *        Created:  10/03/2013 11:38:50 PM
 *       Revision:  none
 *       Compiler:  gcc
 *
 *         Author:  Ji.Zhilong (), [email protected]
 *   Organization:  SJTU
 *
 * =====================================================================================
 */
#include <iostream>
#include <string>
#include <stdexcept>
using namespace std;

enum Color {RED, BLACK};

template <typename Key=int, typename Value=int>
class RBTree
{
public:
  RBTree():head(NULL) {}
  ~RBTree()
  {
	delete head;
  }

  Value get(const Key &key)
  {
	return head->get(key);
  }	

  int size()
  {
	return head->size();
  }

  bool has(const Key &key)
  {
	if (NULL == head)
	  return false;

	try
	{
	  head->get(key);
	}
	catch (out_of_range &e)
	{
	  return false;
	}
	return true;
  }

  void put(const Key &k, const Value &v)
  {
	if (head == NULL)
	  head = new Node(k, v, NULL, NULL, BLACK);
	else {
	  head = put(&head, k, v);
	}
  }

private:
  struct Node
  {
	Key key;
	Value val;
	Node *left;
	Node *right;
	Color color;
	int count;

	Node(const Key &k, const Value &v, Node *l=NULL, Node *r=NULL, Color c=RED, int count=1): 
	  key(k), val(v), left(l), right(r), color(c), count(count){}

	~Node() 
	{
	  if (left != NULL)
		delete left;
	  if (right != NULL)
		delete right;
	}

	int size()
	{
	  if (this == NULL)
		return 0;
	  return count;
	}

	Value get(const Key &key)
	{
	  Node *node = this;
	  while (node != NULL) {
		if (node->key == key) return node->val;
		else if (node->key > key) node = node->left;
		else  node = node->right;
	  }
	  throw out_of_range("wrong key");
	}

	Node *rotate_right()
	{
	  Node *tmp = left;
	  left = tmp->right;
	  tmp->right = this;
	  tmp->color = BLACK;
	  color = RED;
	  update_count();
	  return tmp;
	}

	Node *rotate_left()
	{
	  Node *tmp = right;
	  right = tmp->left;
	  tmp->left = tmp;
	  tmp->color = BLACK;
	  color = RED;
	  update_count();
	  return tmp;
	}

	void flip_color()
	{
	  left->color = BLACK;
	  right->color = BLACK;
	  this->color = RED;
	}

	void update_count()
	{
	  count = left->size() + right->size() + 1;
	}

  };

  Node *put(Node **node, const Key &k, const Value &v)
  {
	Node &n = **node;
	if (*node == NULL)
	  return (new Node(k, v));

	if (k < n.key)
	  n.left = put(&(n.left), k, v);
	else if (k > n.key)
	  n.right = put(&(n.right), k, v);
	else
	  n.val = v;


    if (is_red(n.right) && !is_red(n.left)) 
	  *node = n.rotate_left();

	if (is_red(n.left) && is_red(n.left->left))
	  *node = n.rotate_right();

    if (is_red(n.left) && is_red(n.right))
	  n.flip_color();

	(*node)->update_count();
	return *node;
  }

  static bool is_red(Node *node)
  {
	return (node != NULL && node->color == RED);
  }

  Node *head;
};



int
main()
{
  RBTree<int, string> rb;
  rb.put(1, "one");
  rb.put(-1, "minus one");
  rb.put(3, "three");
  cout << rb.get(3) << endl;
  cout << boolalpha << rb.has(71) << endl;
  cout << rb.size() << endl;
  return(0);
}