// File:     cs2set.h
// Author:   Chuck Stewart
// Purpose:  Solution to Computer Science II, HW 8, Spring 2008.


#ifndef cs2set_h_
#define cs2set_h_

#include <iostream>
#include <utility>
#include <vector>


//  Auxiliary class for nodes in the tree.

template <class T>
class TreeNode {
public:
  TreeNode() : left(0), right(0), parent(0) {}
  TreeNode(const T& init) : value(init), left(0), right(0), parent(0) {}
  T value;
  TreeNode* left;
  TreeNode* right;
  TreeNode* parent;
};


//  Forward declaration of the templated cs2set class so that the
//  iterator class allows it to be named as a friend.

template <class T> class cs2set;


///////////////////////////////////////////////////////////////////
///////////////         Iterator definitions          /////////////
///////////////////////////////////////////////////////////////////

//  These are incomplete because the iterator increment and
//  decrement operators are a bit beyond the scope of this course.

template <class T>
class tree_iterator {
public:
  friend class cs2set<T>;
private:
  TreeNode<T>* ptr_;
public:
  tree_iterator() : ptr_(0) {}
  tree_iterator( TreeNode<T>* p ) : ptr_(p) {}
  tree_iterator( const tree_iterator& old ) : ptr_(old.ptr_) {}
  ~tree_iterator() {}
  
  tree_iterator& operator=( const tree_iterator& old )
  { ptr_ = old.ptr_;  return *this; }
  
  //  operator* gives constant access to the value at the pointer
  const T& operator*() const
  {
    return ptr_->value;
  }

  //  Comparions operators are straightforward
  friend bool operator== ( const tree_iterator& lft, const tree_iterator& rgt )
  { return lft.ptr_ == rgt.ptr_; }
  
  friend bool operator!= ( const tree_iterator& lft, const tree_iterator& rgt )
  { return lft.ptr_ != rgt.ptr_; }


  //--------------------------------------------------
  //  HW 8:  Two ++ and two -- operators
  //--------------------------------------------------

  //  HW 8:  pre-increment
  tree_iterator<T> & operator++( ) 
  { 
    ptr_ = find_next_node( ptr_ );
    return *this;
  }

  //  HW 8:  post-increment is indicated by the dummy int argument
  tree_iterator<T> operator++( int ) 
  {
    tree_iterator<T> temp( *this ); // save an iterator referencing the current node
    ptr_ = find_next_node( ptr_ );
    return temp;   // return what was the current iterator
  }

  //  HW 8:  pre-decrement
  tree_iterator<T> & operator--( ) 
  {
    ptr_ = find_prev_node( ptr_ );
    return *this;
  }

  //  HW 8:  post-decrement is indicated by the dummy int argument
  tree_iterator<T> operator--( int ) 
  {
    tree_iterator<T> temp( *this ); // save an iterator referencing the current node
    ptr_ = find_prev_node( ptr_ );
    return temp;   // return what was the current iterator
  }

private:

  //-----------------------------------------------------------------
  //  HW 8:  find the next node in the tree in an in-order traversal
  //-----------------------------------------------------------------

  TreeNode<T>* find_next_node( TreeNode<T>* p )
  {
    if ( !p ) return p;   // safety check

    //  If p has a non-empty right subtree then the next node is the
    //  left-most node in this right subtree
    if ( p->right )
      {
        p = p->right;
        while ( p->left )
          p = p->left;
        return p;
      }

    //  Otherwise, step back up the tree until the step is back up a
    //  left branch.  The ancestor reached in this way is the next
    //  node.  If the process steps up to the top of the tree, we've
    //  stepped past the last node.
    else 
      {
        TreeNode<T> * q = p->parent;
        while ( q && q->right == p )
          {
            p = q;
            q = q->parent;
          }
        return q;
      }
  }


  //-----------------------------------------------------------------
  //  HW 8:  find the previous node in the tree in an in-order traversal
  //-----------------------------------------------------------------

  TreeNode<T>* find_prev_node( TreeNode<T>* p )
  {
    if ( !p ) return p;

    //  If p has a non-empty left subtree then the pevious node is the
    //  right-most node in this left subtree
    if ( p->left )
      {
        p = p->left;
        while ( p->right )
          p = p->right;
        return p;
      }
   
    //  Otherwise, step back up the tree until the step is back up a
    //  right branch.  The ancestor reached in this way is the next
    //  node.  If the process steps up to the top of the tree, we've
    //  stepped past the first node.
    else 
      {
        TreeNode<T> * q = p->parent;
        while ( q && q->left == p )
          {
            p = q;
            q = q->parent;
          }
        return q;
      }
  }

};



///////////////////////////////////////////////////////////////////
/////////////////       Main cs2set class          ////////////////
///////////////////////////////////////////////////////////////////


template <class T>
class cs2set {
public:
  typedef tree_iterator<T> iterator;

private:
  TreeNode<T>* root_;
  int size_;

public:
  cs2set() : root_(0), size_(0) 
  {}

  cs2set( const cs2set<T>& old ) : size_(old.size_)
  {
    root_ = this -> copy_tree( old.root_ );
  }

  ~cs2set()
  {
    this -> destroy_tree( root_ );
    root_ = 0;
  }

  cs2set& operator=( const cs2set<T>& old )
  {
    if ( old != *this ) {
      this -> destroy_tree( root_ );
      root_ = this ->  copy_tree( (TreeNode<T>*) 0, old.root_ );
      size_ = old.size_;
    }
    return *this;
  }

  int size() const { return size_; }
  
  std::pair< iterator, bool > insert( T const& key_value )
  {
    return insert( key_value, root_ );
  }

  iterator find( const T& key_value )
  {
    return find( key_value, root_ );
  }

  int erase( T const& key_value )
  {
    return erase( key_value, root_ );
  }

  friend std::ostream& operator<< ( std::ostream& ostr, const cs2set<T>& s )
  {
    s.print_in_order( ostr, s.root_ );
    return ostr;
  }

  void print_as_sideways_tree( std::ostream& ostr ) const
  {
    print_as_sideways_tree( ostr, root_, 0 );
  }

  iterator begin() const
  { 
    if ( ! root_ )
      return iterator(0);
    else
      {
        TreeNode<T>* p = root_;
        while ( p->left ) p = p->left;
        return iterator(p);
      }
  }

  iterator end() const { return iterator( 0 ); }


  //--------------------------------------------------
  //  Here are the functions to implement for HW 8
  //--------------------------------------------------

  //  HW 8:  Find the height of the tree
  int height() const 
  { 
    return compute_height( root_ );
  }

  //  HW 8: Find the average depth of the nodes in the tree
  float average_depth() const
  { 
    if ( !root_ ) return -1;
    int sum_depths = compute_accumulated_depths( root_, 0 );
    return float(sum_depths) / float(size_);
  }

  //  HW 8:  Rebalance the tree
  void rebalance() 
  {
    std::vector< TreeNode<T>* > node_ptrs;
    copy_nodes_into_vector( root_, node_ptrs );
    root_ = redistribute_nodes( node_ptrs, 0, node_ptrs.size()-1 );
    if ( root_ ) root_->parent = 0;
  }

  //  HW 8:  Determine if the tree is balanced.
  bool is_balanced() const 
  { 
    int num_in_subtree;
    return is_balanced( root_, num_in_subtree ); 
  }

  //  HW 8:  Output the values in level order
  void level_order( std::ostream& ostr ) const 
  {
    std::vector< std::vector< TreeNode<T>* > > ptrs_at_levels;
    copy_ptrs_to_levels( root_, 0, ptrs_at_levels );
    for ( unsigned int i=0; i<ptrs_at_levels.size(); ++i )
      {
        ostr << "Level " << i << ": ";
        for ( unsigned int j=0; j<ptrs_at_levels[i].size(); ++j )
          ostr << ' ' << ptrs_at_levels[i][j]->value;
        ostr << '\n';
      }
  }
  
  //  HW 8:  Apply operator==
  friend bool operator==( const cs2set<T>& s, const cs2set<T>& t ) 
  { 
    if ( s.size_ != t.size_ ) return false;
    std::vector< TreeNode<T>* > s_node_ptrs, t_node_ptrs;
    s.copy_nodes_into_vector( s.root_, s_node_ptrs );    
    t.copy_nodes_into_vector( t.root_, t_node_ptrs );    
    assert( s_node_ptrs.size() == t_node_ptrs.size() );
    for ( unsigned int i=0; i < s_node_ptrs.size(); ++i )
      if ( s_node_ptrs[i]->value < t_node_ptrs[i]->value ||
           t_node_ptrs[i]->value < s_node_ptrs[i]->value )
        return false;
    return true;
  }

  //  HW 8:  lower_bound functions, similar to std::set<T>
  iterator lower_bound( const T& x ) const 
  {
    // Keep a pointer to the current lower bound based on how far
    // we've looked through the tree.
    TreeNode<T> * lower_ptr = 0;

    //  Go down the tree.  Whenever a new node is found that is better
    //  than the previous lower bound, update the lower ptr and move
    //  left through the tree.
    TreeNode<T>* p = root_;
    while ( p )
      {
        if ( ! ( p->value < x ) )
          {
            lower_ptr = p;
            p = p->left;
          }
        else
          p = p->right;
      }
    return iterator(lower_ptr);
  }

  //  HW 8:  upper_bound functions, similar to std::set<T>
  iterator upper_bound( const T& x ) const
  { 

    // Keep a pointer to the current upper bound based on how far
    // we've looked through the tree.
    TreeNode<T> * upper_ptr = 0;

    //  Go down the tree.  Whenever a new node is found that is better
    //  than the previous upper bound, update the upper ptr and move
    //  left through the tree.
    TreeNode<T>* p = root_;
    while ( p )
      {
        if ( x < p->value )
          {
            upper_ptr = p;
            p = p->left;
          }
        else
          p = p->right;
      }
    return iterator(upper_ptr);
  }
  
  //  HW 8:  Erase the node referenced by the iterator.
  void erase( iterator p ) 
  {
    //  Make use of the previous erase function, but step up to the
    //  parent node so that its child pointer can be adjusted
    TreeNode<T>* q = p.ptr_;
    if ( q->parent && q->parent->left == q )
      {
        //  pass parent's left child pointer so that is it adjusted correctly
        erase( q->value, q->parent->left );
      }
    else if ( q->parent && q->parent->right == q )
      {
        //  pass parent's right child pointer so that is it adjusted correctly
        erase( q->value, q->parent->right );
      }
    else
      erase( q->value, root_ );   //  the value being removed is at the root
  }

private:

  TreeNode<T>*  copy_tree( TreeNode<T>* old_root )
  {
    //  Base case is an empty (sub)tree
    if ( !old_root ) return 0;

    //  Create the new node
    TreeNode<T>* new_root = new TreeNode<T>( old_root->value );

    //  Recursively build the left subtree.  If it is non-null, set
    //  the parent pointer of the left child to point back to the new
    //  root.
    new_root->left = copy_tree( old_root->left );
    if ( new_root->left ) new_root->left->parent = new_root;

    //  Recursively build the right subtree.  If it is non-null, set
    //  the parent pointer of the right child to point back to the new
    //  root.
    new_root->right = copy_tree( old_root->right );
    if ( new_root->right ) new_root->right->parent = new_root;

    return new_root;
  }

  void destroy_tree( TreeNode<T>* p )
  {
    if ( p )
      {
	destroy_tree(p->left);
	destroy_tree(p->right);
	delete p;
	p = 0;
      }
  }

  void print_in_order( std::ostream& ostr, const TreeNode<T>* p ) const
  {
    if ( p )
      {
	print_in_order( ostr, p->left );
	ostr << p->value << "\n";
	print_in_order( ostr, p->right );
      }
  }

  std::pair<iterator,bool> insert( const T& key_value, TreeNode<T>*& p )
  {
    if ( !p )
      {
	p = new TreeNode<T>( key_value );
	this -> size_ ++ ;
	return std::pair<iterator,bool>( iterator(p), true );
      }

    //  The rest is a non-recursive version of insert, which moves
    //  left and right down the tree until the location is found to
    //  insert.  The loop has no termination condition because
    //  termination is handled by the return statements in the if
    //  conditions.  
    TreeNode<T>* q = p;  // Work with q because p was passed by reference.
    while ( true )
      {
        //  If the value should be in the the left subtree
        if ( key_value < q->value )
          {
            if ( !q->left )
              {
                //  Reached the bottom of tree so insert a new node
                //  and set its parent points
                q->left = new TreeNode<T>( key_value );
                q->left->parent = q;
                size_ ++ ;
                return std::pair<iterator,bool>( iterator(q->left), true );
              }
            else
              q = q->left;
          }

        //  If the value should be in the right subtree.
        else if ( key_value > q->value )
          {
            if ( !q->right )
              {
                //  Reached the bottom of tree so insert a new node
                //  and set its parent points
                q->right = new TreeNode<T>( key_value );
                q->right->parent = q;
                size_ ++ ;
                return std::pair<iterator,bool>( iterator(q->right), true );
              }
            else
              q = q->right;
          }

        //  The value is at the current location, so we're done.
        else
          return std::pair<iterator,bool>( iterator(q), false );
      }
  }

  iterator find( const T& key_value, TreeNode<T>* p )
  {
    if ( !p )
      return this -> end();
    else if ( key_value < p->value )
      return find( key_value, p->left );
    else if ( key_value > p->value )
      return find( key_value, p->right );
    else
      return iterator( p );
  }

  int erase( T const& key_value, TreeNode<T>* & root )
  {
    if ( !root ) return 0;
    
    if ( key_value < root->value )
      return erase( key_value, root->left );
    else if ( key_value > root->value )
      return erase( key_value, root->right );

    if ( !root->left )
      {
	TreeNode<T> * temp = root;
	root = root->right;

        //  Fix the parent pointer of the promoted node if this node exists.
        if ( root ) root->parent = temp->parent;
	
	size_ --;
	delete temp;
	return 1;
      }
    else if ( !root->right )
      {
	TreeNode<T> * temp = root;
	root = root->left;

        //  Fix the parent pointer of the promoted node if this node exists.
        if ( root ) root->parent = temp->parent;

	size_ --;
	delete temp;
	return 1;
      }

    //  Two children  in-order successor
    TreeNode<T>* leftmost = root->right;
    while ( leftmost->left ) leftmost = leftmost->left;
    root->value = leftmost->value;
    //size_ --;
	return erase( root->value, root->right );
  }


  void print_as_sideways_tree( std::ostream& ostr, const TreeNode<T>* p,
                               int depth  ) const
  {
    if ( p ) 
      {
        print_as_sideways_tree( ostr, p->right, depth+1 );
        for ( int i=0; i<depth; ++i ) ostr << "    ";
        ostr << p->value << "\n";
        print_as_sideways_tree( ostr, p->left, depth+1 );
      }
  }


  //------------------------------------------------------------
  //  HW 8:  extra private functions
  //------------------------------------------------------------

  int compute_height( TreeNode<T> * root ) const
  {
    if ( !root ) 
      return -1;
    else 
      return 1 + std::max( compute_height(root->left), compute_height(root->right) );
  }

  int compute_accumulated_depths( TreeNode<T> * root, 
                                  int current_depth ) const
  {
    if ( !root ) 
      return 0;   // base case
    else
      //  accumulate the left and right depths by adding one to the
      //  current depth and recursing
      return 
        compute_accumulated_depths( root->left, current_depth+1 ) 
        + compute_accumulated_depths( root->right, current_depth+1 )
        + current_depth;
  }

  //  HW 8:  Rebalance the tree
  void copy_nodes_into_vector( TreeNode<T>* root,
                               std::vector< TreeNode<T>* > & node_ptrs ) const
  {
    if ( !root ) return;
    //  Copy the nodes in the a vector using a standard in order
    //  traversal.
    copy_nodes_into_vector( root->left, node_ptrs );
    node_ptrs.push_back( root );
    copy_nodes_into_vector( root->right, node_ptrs );
  }    

  TreeNode<T>* redistribute_nodes( std::vector< TreeNode<T>* > & node_ptrs,
                                   int low, 
                                   int high )
  {
    if ( high < low ) return (TreeNode<T>*) 0;

    //  Split the tree, making the mid value the root and recursively
    //  building the left and right subtrees.
    int mid = (low+high+1)/2;
    TreeNode<T>* root = node_ptrs[mid];
    root->left = redistribute_nodes( node_ptrs, low, mid-1 );
    root->right = redistribute_nodes( node_ptrs, mid+1, high );

    //  Assign the parent pointers of non-null children
    if ( root->left ) root->left->parent = root;
    if ( root->right ) root->right->parent = root;
    return root; 
  }


  //  HW 8:  Determine if the tree is balanced.
  bool is_balanced( TreeNode<T>* root,
                    int & num_in_subtree ) const
  {
    //  Keep track of the number of the nodes in the subtree.
    if ( !root )
      {
        num_in_subtree = 0;
        return true;
      }
    
    //  Check the left and right balance and count nodes in the
    //  subtrees.  If one of subtrees is unbalanced, the whole thing
    //  is unbalanced.
    int num_in_left, num_in_right;
    if ( !is_balanced( root->left, num_in_left ) ) return false;
    if ( !is_balanced( root->right, num_in_right ) ) return false;
    if ( std::abs( num_in_left - num_in_right ) > 1 ) return false;

    num_in_subtree = num_in_left + num_in_right + 1;
    return true;
  }


  //  HW 8:  Output the values in level order
  void copy_ptrs_to_levels( TreeNode<T>* root, 
                            int current_level,
                            std::vector< std::vector< TreeNode<T>* > > & ptrs_at_levels ) const
  {
    if ( !root ) return;

    //  If we have reached a level where there is no current vector,
    //  add an empty vector.
    if ( current_level >= int(ptrs_at_levels.size()) )
      {
        std::vector< TreeNode<T>* > empty;
        ptrs_at_levels.push_back( empty );
      }
    
    //  Add the pointer to the current level and recurse to the
    //  children.
    ptrs_at_levels[current_level].push_back( root );
    copy_ptrs_to_levels( root->left, current_level+1, ptrs_at_levels );
    copy_ptrs_to_levels( root->right, current_level+1, ptrs_at_levels );
  }
};



#endif