// © Copyright 1997. Joseph Bergin. All rights reserved.

#ifndef REDBLACKTREE_H
#define REDBLACKTREE_H

// Implements a balanced binary search tree with operator< as the comparison
// operator.  No leaf has more than twice the height of any other.  If a<b are
// both in the tree then a is to the left of b OR a is a left descendant of b OR
// b is a right descendant of a.  

// Logically a 2-3-4 tree with all leaves at exactly the same height.  The 234 
// nodes are built from red and black binary nodes as follows. 
// A binary node with two black children is logically a 2-node.
// A binary node with one red and one black child is logically a three-node.
// A binary node with red nodes for both children is logically a four node.
// The trailer node in each tree has no color.

// T requires a default constructor and destructor (unless it is a pointer type),
// operator<, operator==, operator=.  

// reference: L.J. Guibas & R. Sedgewick, "A Dichromatic Framework for Balanced Trees," 
// Proceedings of the Nineteenth Annual IEEE Symposium on Foundations of Computer 
// Science (1978, 8-21. 

// The insert algorithm and the rotate and split functions are adapted from Sedgewick, 
// "Algorithms" Second Edition, Addison Wesley, 1988.  The erase algorithm was adapted
// from the HP STL implementation.  

template <class T>
class RedBlackTree
{	protected:
		struct node;
		typedef node* nodeptr;
		static nodeptr Z; // used to terminate all chains
		struct node
		{	enum color{red, black,none};
			T _data;
			color _color;
			nodeptr _left;
			nodeptr _right;
			nodeptr _parent;
			node(T data = T(), color c = none)
			:	_color(c),
				_left(Z),
				_right(Z),
				_parent(Z),
				_data(data)
			{
			}
		};
		nodeptr _root;
			// This node is a dummy "header".  The actual tree is to its right. 
		nodeptr _trailer;
			// This node is the root of an empty tree and is the right child of
			// the rightmost logical entry in other trees.  It represents a 
			// past-the-end location. 
		long	_nodeCount;
		
	public:
		typedef T value_type;
		typedef T* pointer;
		typedef T& reference;
		
		RedBlackTree():_nodeCount(0),_root(new node()),_trailer(new node())
		{	if(Z == NULL)
			{	Z = new node();
				Z->_left = Z;
				Z->_right = Z;
				Z->_parent = Z;
				_root->_left = Z;
				_root->_parent = Z;
				_trailer->_left = Z;
				_trailer->_right = Z;
				_trailer->_data = -999; // debug only
			}
			_root->_right = _trailer;
			_trailer->_parent = _root;
		}
		
		bool empty()const{return _root->_right == _trailer;}
		
		class iterator // Preorder iterator.  Inorder would be far superior. 
		{	public:
				iterator(RedBlackTree<T> & t):_here(t._root->_right),_tree(t){};
				bool operator==(const RedBlackTree<T>::iterator& i)const
				{	return _here == i._here;
				}
				
				T& operator*(){return _here->_data;}
				
				iterator& operator++()
				{	if( _here == _tree._trailer){return *this;}
					if( _here->_left != Z) _here = _here->_left;
					else if(_here->_right != Z) _here = _here->_right;
					else
					{	nodeptr old;
						do
						{	old = _here;
							_here = _here->_parent;
						}	while(_here->_right == Z || _here->_right == old);
						_here = _here->_right;
					}
					return *this;
				}

				iterator& operator--()
				{	if(_here->_parent == _tree._root){return *this;} // potential bug here
					if(_here == _here->_parent->_left) _here = _here->_parent;
					else
					{	_here =  _here->_parent;
						while(_here->_left != Z)
						{	_here = _here->_left;
							while(_here->_right != Z) // Will never be _trailer.
								_here = _here->_right;
						}
					}
					return *this;
				}
								
								
			protected:
				nodeptr _here;
				RedBlackTree<T>& _tree;
				friend class RedBlackTree<T>;
		};

		iterator begin()const
		{	return iterator(*this);
		}
		
		iterator end()const
		{	iterator result(*this);
			result._here = _trailer;
			return result;
		}
		
	protected:
	
		enum rotation{left, right};
		
		nodeptr rotate(bool rightRotation, nodeptr where)
		// where is the the point of rotation.  
		// If rightRotation is true we will rotate right, otherwise
		// the rotation will be left.  
		// We must always rotate to the right of the _root, however. 
		{	nodeptr parent = where->_parent, result;
			bool leftChild = (where == parent->_left);
			if(rightRotation)
			{	result = where->_left;
				where->_left = result->_right;
				where->_left->_parent = where;
				result->_right = where;
			}
			else
			{	result = where->_right;
				where->_right = result->_left; 
				where->_right->_parent = where;
				result->_left = where;
			}
			where->_parent = result;
			if(leftChild)
			{	parent->_left =  result;
			} 
			else 
			{	parent->_right = result;
			}
			result->_parent = parent;
			return result;
		}
		
		nodeptr split(const T& val, nodeptr grand, nodeptr parent, nodeptr here)
		// Split a 4-node into 2 2-nodes.
		{	here->_color = node::red;
			here->_left->_color = node::black;
			here->_right->_color = node::black;
			if(parent->_color == node::red)
			{	grand->_color = node::red;
				if((val < grand->_data) != (val < parent->_data)) 
					parent = rotate(val < parent->_data, parent);
				here = rotate(val < grand->_data, grand);
				here->_color = node::black;
			}
			_root->_right->_color = node::black;
			return here;
		}
		
	public:	
		iterator insert(const T& t)
		// Returns an iterator to the inserted item.
		{	_nodeCount++;
			iterator result(*this);
			if(empty())
			{	nodeptr newRoot = new node(t, node::black);
				newRoot->_right = _trailer;
				_trailer->_parent = newRoot;
				_root->_right = newRoot;
				newRoot->_parent = _root;
				result._here = newRoot;
			}
			else
			{	nodeptr grand = _root, parent = _root, here = _root;
				do
				{	grand = parent; parent = here;
					if( here != _root && t < here->_data) here = here->_left; 
					else here = here->_right;
					// Split any 4-nodes you encounter
					if(here->_left->_color == node::red && here->_right->_color == node::red)
						here = split(t, grand, parent, here);
				}while(here != _trailer && here != Z);
				nodeptr temp = here;
				here = new node(t,node::black);
				here->_right = temp;  // Preserve the trailer.
				if(temp == _trailer)_trailer->_parent = here;
				here->_parent = parent;
				if(t < parent->_data)parent->_left = here; else parent->_right = here;
				result._here = here;
				here = split(t, grand, parent, here);
			}
			return result;
		}
		
		iterator find(const T& t) const
		// Return an iterator to t's location if present.  
		// Returns a past the end location otherwise. 
		{	iterator result(*this);
			nodeptr here = _root->_right;
			while(here != Z && here->_data != t)
				if(here->_data < t) here = here->_right; else here = here->_left;
			if(here == Z) here = _trailer;
			result._here = here;
			return result;
		}

		iterator findNext(iterator from)
		// finds the next occurence of *from or returns end()
		{	iterator result = from;
			iterator top = find(*from);
			// top is the topmost occurence of the desired value.  If we reach it we have
			// returned all of the values.  
			do
			{	++result;
			}while(result._here != _trailer && result != top && *result != *from);
			if(result == top) result._here = _trailer;
			return result;
		}
		
		void erase(iterator where)
		// Erase the item pointed to by where. 
		// Works by finding a node near a leaf to swap with the node at where.  It then removes
		//  this node and rebalances the tree by working upwards from that leaf. 
		{	nodeptr here = where._here;
			nodeptr y = here; // This will be the location actually removed. 
			nodeptr replacement; // Replaces the moved node. 
			if(y->_left == Z) replacement = y->_right; // Removing here itself.
			else if(y->_right == Z) replacement = y->_left; // Removing here itself.
			else
			{	y = y->_right;
				while(y->_left != Z) y = y->_left;// Smallest value in right tree. Removing y's location.
				replacement = y->_right;
			}
			if(y != here) // Relink y into here's location
			{	here->_left->_parent = y;
				y->_left = here->_left;
				if(y != here->_right)
				{	replacement->_parent = y->_parent;
					y->_parent->_left = replacement;
					y->_right = here->_right;
					here->_right->_parent = y;
				}
				else
					replacement->_parent = y;
				if(here->_parent->_right == here)
					here->_parent->_right = y;
				else  
					here->_parent->_left = y;
				y->_parent = here->_parent;
				::swap(y->_color, here->_color);
				y = here; // we will delete this node later. 
			}
			else //  y == here. Remove y. 
			{	replacement->_parent = y->_parent;
				if(here->_parent->_right == here)
					here->_parent->_right = replacement;
				else  
					here->_parent->_left = replacement;
			}
			// Rebalance.  Only if we are shortening the equivalent 2-3-4 tree is this
			// necessary. We need to either lengthen this part of the tree by one to 
			// compensate for the removed node OR shorten the rest of the tree by one.  
			// We don't know yet which is easiest.  
			if(y->_color == node::black) // Removing a black y node will shorten the tree. 
			{	nodeptr x = replacement; // x is the adjustment location. We are too short here. 
				while( x != _root->_right && x->_color == node::black)
					if(x == x->_parent->_left) // x is a left child. 
					{	nodeptr sib = x->_parent->_right; // sibling of x.  
						if (sib->_color == node::red) // Then its parent and children are black.
						{	sib->_color = node::black;		//	|	Leaves logical heights unaffected.
							x->_parent->_color = node::red; // 	| 	Logical height = #black nodes  
							rotate(left, x->_parent);	 	//  |   to root.
							sib = x->_parent->_right;	// Old sib's left child. 
						} // sib is now black. x and Sib are at same height. Sib is still longer than x.   
						if(sib->_left->_color == node::black && sib->_right->_color == node::black)
						{	sib->_color = node::red; // Shorten sib's chain. Balance parent. 
							x = x->_parent; // continue up the tree;
						}
						else // Sib is still longer, but we can fix it.  Note: one child of sib is red here.
						{	if(sib->_right->_color == node::black) // Then left is red. & its children black
							{	sib->_left->_color = node::black;	// | Leaves logicla heights unaffected
								sib->_color = node::red;			// | but gives x a black sib with a black
								rotate(right, sib);					// | left child.
								sib = x->_parent->_right;
							}// sib is black and it's left is black also. 
							sib->_color = x->_parent->_color;
							x->_parent->_color = node::black; 	// This will lower x if we rotate 
							sib->_right->_color = node::black; 	// and preserve color of new local root.
							rotate(left, x->_parent); 	// lower's x so
							break; 						// all done here
						}
					}
					else // x is a right child instead (mirror reflection of then clause).
					{	nodeptr sib = x->_parent->_left;
						if (sib->_color == node::red)
						{	sib->_color = node::black;
							x->_parent->_color = node::red;
							rotate(right, x->_parent);
							sib = x->_parent->_left;
						} // sib is now black
						if(sib->_right->_color == node::black && sib->_left->_color == node::black)
						{	sib->_color = node::red;
							x = x->_parent; // continue up the tree;
						}
						else
						{	if(sib->_left->_color == node::black)
							{	sib->_right->_color = node::black;
								sib->_color = node::red;
								rotate(left, sib);
								sib = x->_parent->_left;
							}
							sib->_color = x->_parent->_color;
							x->_parent->_color = node::black;
							sib->_left->_color = node::black;
							rotate(right, x->_parent);
							break; // all done here
						}

					}
				x->_color = node::black;
			}
			delete y;
			_nodeCount--;
		}
		
		void tabs(int n)
		{ for(int i = 0; i<n; ++i) cout<<'	';
		}
		
		void dumpNode(nodeptr t, int n)
		{	if(t->_left != Z) dumpNode(t->_left, n+1);
			tabs(n); cout<< t->_data<<(t->_color==node::red?" red":" black")<<endl;
			if(t->_right != Z && t->_right != _trailer) dumpNode(t->_right, n+1);	
		}
		
		void dump()
		{	dumpNode(_root->_right, 0);
		}
		
	friend class RedBlackTree<T>::iterator;
	friend class RedBlackTree<T>::node;
};

template <class T> RedBlackTree<T>::nodeptr RedBlackTree<T>::Z = NULL;

#endif