//
//  File:  BSTree.h
//  Author:  Chuck Stewart
//
//  Purpose:  Class implementing either a balanced or unbalanced
//    binary search tree.  The balancing creates an AVL tree.   This
//    differs from the author's implementation in several ways.  One
//    is that it implements non-lazy deletion.  A second is that
//    unbalanced and balanced trees are combined in the same
//    implementation to show their similarity.  A third is the
//    addition of a fairly simple, sideways printing function to aid
//    in visualizing the tree.
//

#include <iostream>
using namespace std;

//  Utility function.
inline int Max(int a, int b) { return a>b ? a : b; }


//  Simple template for nodes in the binary search tree.   All
//  variables are public.
//
template <class Comparable>
class BSNode {
public:
  Comparable element;
  BSNode<Comparable> *left;
  BSNode<Comparable> *right;
  int height;
  BSNode( Comparable E = 0, BSNode<Comparable> *L = 0,
	    BSNode<Comparable> *R = 0, int H = 0 )
     : element ( E ), left( L ), right( R ), height( H ) { }
};


//  The BSTree class.  If the parameter use_balancing is set true, an
//  AVL tree is realized.

template <class Comparable>
class BSTree {
public:
  BSTree( bool use_balancing=true ) 
    :  root_(0), balanced_(use_balancing)  {}
  BSTree( BSTree<Comparable> & old ) : balanced_(old.balanced_) 
    { root_ = Copy( old->root_ ); }

  // Destructor.
  ~BSTree( )  { Make_Empty(root_); root_ = 0; }
  
  // Assignment operator.
  const BSTree<Comparable> & operator = ( const BSTree<Comparable> & old );

  // Key member functions.  Most are simple drivers.
  void Make_Empty( )
    { Make_Empty ( root_ ); root_ = 0; }
  void Print_Tree( ) const
    { Print_Tree( root_, 0 ); }
  bool Find( const Comparable & X )
    { return Find( X, root_ ); }
  void Insert( const Comparable & X )
    { Insert( X, root_ ); }
  void Remove( const Comparable & X )
    { Remove( X, root_ ); }
  int Height( )
    { return (root_ == 0) ? -1 : root_->height; }

private:
  //  The private member functions do the real work.
  void Make_Empty ( BSNode<Comparable> * & t );
  BSNode<Comparable> * Copy( const BSNode<Comparable> *t );
  void Insert( const Comparable & X, BSNode<Comparable> * & t );
  void Remove( const Comparable & X, BSNode<Comparable> * & t );
  void Print_Tree( BSNode<Comparable> *  t, int depth ) const;
  BSNode<Comparable> *
    Find( const Comparable & X, BSNode<Comparable> * t ) const;
  BSNode<Comparable> * 
    Find_Min( BSNode<Comparable> * t ) const;
  int Node_Ht( BSNode<Comparable> *p ) const;
  void Calculate_Height( BSNode<Comparable> *p ) const;

  void Balance_Left( BSNode<Comparable>* & t ) const;
  void Balance_Right( BSNode<Comparable>* & t ) const;
  void Rotate_With_Left_Child( BSNode<Comparable>* & t ) const;
  void Rotate_With_Right_Child( BSNode<Comparable>* & t ) const;

private:
  //  The root and the balancing flag are the only member variables.
  BSNode<Comparable> *root_;
  bool balanced_;
};



template <class Comparable>
const BSTree<Comparable> &
BSTree<Comparable>::
operator = ( const BSTree<Comparable> & old )
{
  if( this != &old ) {
    Make_Empty( root_ );
    root_ = Copy( old.root_ );
    balanced_ = old.balanced_;
  }
  return *this;
}


template <class Comparable>
BSNode<Comparable> *
BSTree<Comparable>::
Copy( const BSNode<Comparable> *t )
{
  if( t != 0 )
    return
      new BSNode<Comparable> ( t->element, Copy( t->left ),
				Copy( t->right ), t->height );
  else
    return 0;
}


template <class Comparable>
void
BSTree<Comparable>::
Make_Empty( BSNode<Comparable> * & t )
{
  if( t != 0 ) {
    Make_Empty( t->left );
    Make_Empty( t->right );
    delete t;
    t = 0;
  }
}


//  None of the Find functions require recursion.

template <class Comparable>
BSNode<Comparable>*
BSTree<Comparable>::
Find( const Comparable & X, BSNode<Comparable> * t ) const
{
  while ( t != 0 ) {
    if ( X < t->element )
      t = t->left;
    else if ( X > t->element )
      t = t->right;
    else
      return t;
  }
  return 0;
}


template <class Comparable>
BSNode<Comparable>*
BSTree<Comparable>::
Find_Min( BSNode<Comparable> *t ) const
{
  while ( t != 0 ) {
    if ( t->left == 0 ) break;
    t = t->left;
  }
  return t;   // returns 0 if the tree is empty
}


// The next two functions are height calculation utilities.

template <class Comparable>
int 
BSTree<Comparable>::
Node_Ht( BSNode<Comparable> *p ) const
{ 
  if ( p )
    return p->height;  
  else
    return -1;   // null pointers are height -1
}


template <class Comparable>
void 
BSTree<Comparable>::
Calculate_Height( BSNode<Comparable> *p ) const
{ 
  p->height = 1 + Max( Node_Ht( p->left ), Node_Ht( p->right ) ); 
}



//
//  Insert a value into the tree.  After insertion the tree is
//  balanced as necessary (if the balance flag is set).  In any case,
//  the heights of the nodes are recalculated as the recursive
//  unwinds.   As a slight aside:  the actual work of rebalancing will
//  occur at most once for each insertion.
//
template <class Comparable>
void
BSTree<Comparable>::
Insert( const Comparable & X, BSNode<Comparable> * & t )
{
  if( t == 0 ) {    // Create a one node tree.
    t = new BSNode<Comparable>( X );
  }
  else if( X < t->element ) {
    Insert( X, t->left );
    if ( balanced_ )
      Balance_Left( t );
    else
      Calculate_Height( t );
  }
  else if( X > t->element ) {
    Insert( X, t->right );
    if ( balanced_ )
      Balance_Right( t );
    else
      Calculate_Height( t );
  }
  // Else X is in the tree already, so we'll do nothing.
}



//  Remove a value from the tree.  This does not use lazy deletion.
//  Like insert, balancing occurs after removal if the balancing flag
//  is set.
//
template <class Comparable>
void
BSTree<Comparable>::
Remove( const Comparable & X, BSNode<Comparable> * & t )
{
  BSNode<Comparable> * tmp;

  if ( t == 0 ) {
    cerr << "element not found in \"Remove\" " << "\n";
  }
  else if ( X < t->element ) {  // Look in left subtree
    Remove( X, t->left );
    if( balanced_ )
      Balance_Right( t );  // Right subtree may now be too tall
    else
      Calculate_Height( t );
  }
  else if ( X > t->element ) {  // Look in left subtree
    Remove( X, t->right );
    if ( balanced_ )
      Balance_Left( t );  // Left subtree may now be too tall
    else
      Calculate_Height( t );
  }

  // element is found.  Remove now depends on whether the node
  // containing X has 0, 1, or 2 children
  //
  else if ( t->left != 0 && t->right != 0 ) {
    tmp = Find_Min( t->right );
    t->element = tmp->element;
    Remove( t->element, t->right );
    if ( balanced_ )
      Balance_Left( t );  // left subtree may now be too tall
    else
      Calculate_Height( t );
  }
  else {
    tmp = t;
    // No height recalculation needed here.
    if ( t->left == 0 )
      t = t->right; 
    else
      t = t->left;  
    delete tmp;
  }
}



template <class Comparable>
void
BSTree<Comparable>::
Print_Tree( BSNode<Comparable> *t, int depth ) const
{
  if( t != 0 ) {
    Print_Tree( t->right, depth+1 );
    for ( int i=0; i<depth; i++ ) cout << "   ";
    cout << t->element << "," << t->height << endl ;
    Print_Tree( t->left, depth+1 );
  }
}



///////////  Balancing functions  /////////////////////

//  For a given node t, check whether the left subtree has gotten too
//  large relative to the right subtree.  This could occur if a node
//  has been added to the left subtree or deleted from the right.  If
//  the left subtree has gotten too big, then either do a single or a
//  double rotation.  If the left subtree has not gotten too big, then
//  just calculate the height of t.
//
template <class Comparable>
void
BSTree<Comparable>::
Balance_Left( BSNode<Comparable> * & t ) const
{
  if ( Node_Ht( t->left ) - Node_Ht( t->right ) == 2 ) {
    if ( Node_Ht( t->left->right ) > Node_Ht( t->left->left ) ) {
      //  Double rotation from the left.
#ifdef DEMO
      cout << "Double rotation from left.\n";
#endif
      Rotate_With_Right_Child( t->left );
      Rotate_With_Left_Child( t );
    }
    else {
      //  Single rotation from the left.
#ifdef DEMO
      cout << "Single rotation from the left.\n";
#endif
      Rotate_With_Left_Child( t );
    }
  }
  else {
    Calculate_Height( t );
  }
}


//
//  For a given node T, check whether the right subtree has gotten too
//  large relative to the left subtree.  This could occur if a node
//  has been added to the right subtree or deleted from the left.  If
//  the right subtree has gotten too big, then either do a single or a
//  double rotation.  If the right subtree has not gotten too big, then
//  just calculate the height of t.
//
template <class Comparable>
void
BSTree<Comparable>::
Balance_Right( BSNode<Comparable> * & t ) const
{
  if ( Node_Ht( t->right ) - Node_Ht( t->left ) == 2 ) {
    if ( Node_Ht( t->right->left ) > Node_Ht( t->right->right ) ) {
#ifdef DEMO
      cout << "Double rotation from right.\n";
#endif
      Rotate_With_Left_Child( t->right );
      Rotate_With_Right_Child( t );
    }
    else {
#ifdef DEMO
      cout << "Single rotation from the right.\n";
#endif
      Rotate_With_Right_Child( t );
    }
  }
  else {
    Calculate_Height( t );
  }
}


//  The following two functions do the actual work of rotations.

template <class Comparable>
void
BSTree<Comparable>::
Rotate_With_Left_Child( BSNode<Comparable>* & t ) const
{
  BSNode<Comparable> * k1 = t->left;
  t->left = k1->right;
  k1->right = t;
  t->height = 1 + Max( Node_Ht( t->left ), Node_Ht( t->right ) );
  k1->height = 1 + Max( Node_Ht( k1->left ), Node_Ht( k1->right ) );
  t = k1;
}      


template <class Comparable>
void
BSTree<Comparable>::
Rotate_With_Right_Child( BSNode<Comparable>* & t ) const
{
  BSNode<Comparable> * k1 = t->right;
  t->right = k1->left;
  k1->left = t;
  t->height = 1 + Max( Node_Ht( t->left ), Node_Ht( t->right ) );
  k1->height = 1 + Max( Node_Ht( k1->left ), Node_Ht( k1->right ) );
  t = k1;
}