Data Structures, Algorithms, & Applications in Java
Chapter 15, Exercise 19

First we must decide how we are going to incorporate the field leftSize. We can either require the user to include this field as one of the fields of the elements he/she will insert into the search tree, or we can define a new type IndexedElement with a leftSize field and a field element of type Object. In our implementation of IndexedBinarySearchTree we have gone the second route.

The code for the class IndexedElement is given below. 
static class IndexedElement
{
   // data members
   Object element;        // element in node
   int leftSize;          // number of nodes in left subtree

   // constructors
   IndexedElement(int theLeftSize, Object theElement)
   {
      leftSize = theLeftSize;
      element = theElement;
   }

   IndexedElement(Object theElement)
      {element = theElement;}

   /** @return equivalent string suitable for output */
   public String toString()
      {return element.toString();}
}


The class IndexedBinarySearchTree is derived from BinarySearchTree. The method ascend as well as the constructor are simply inherited from the base class.

The code for the method get(theKey) invokes BinarySearchTree.get(theKey), which returns an element of type IndexedElement. The element field of this returned object is extracted and returned as the result. The code is given below. 
/** @return element whose key is theKey
  * @return null if there is no element with key theKey */
public Object get(Object theKey)
{
   Object theElement = super.get(theKey);
   if (theElement == null)
      // no matching element
      return null;
   else
      return ((IndexedElement) theElement).element;
}


To facilitate the implementation of the method get(theIndex), we first develop the method getDataElement(theIndex), which returns the element field (this is of type Data) of the node whose index is theIndex. 
/** @return Data object of element whose index is theIndex
  * @return null if there is no element with index theIndex */
Data getDataElement(int theIndex)
{
   BinaryTreeNode p = root;
   while (p != null)
   {
      IndexedElement q = (IndexedElement) ((Data) p.element).element;
      if (theIndex < q.leftSize)
         // desired element is in left subtree
         p = p.leftChild;
      else
         if (theIndex > q.leftSize)
         {// desired element is in left subtree
   	        theIndex -= q.leftSize + 1;
             p = p.rightChild;
         }
         else 
            // found desired element
            return (Data) p.element; 
   }

   // no element with index equal to theIndex
   return null;
}


Now, we use the method getDataElement to arrive at the code for get(theIndex). 
/** @return element whose index is theIndex
  * @return null if there is no element with index theIndex */
public Object get(int theIndex)
{
   Data theData = getDataElement(theIndex);
   if (theData == null)
      // no element with index theIndex
      return null;
   else
      return ((IndexedElement) theData.element).element;
}


An insert may be performed in two phases. First, an insert is done assuming we have an ordinary binary search tree (not an indexed one). If the tree did not already contain an element with the given key, then another pass is made from the root to the newly inserted node adjusting the value of leftSize as needed. If we use BinarySearchTree.put for the first pass, we must determine the size of the search tree before and after the invocation of BinarySearchTree.put so that we can determine whether or not the second pass is to be made. An alternative to this is to just replicate the code of BinarySearchTree.put into IndexedBinarySearchTree.put. We take this alternative route below. 
/** insert an element with the specified key
  * overwrite old element if there is already an
  * element with the given key
  * @return old element (if any) with key = theKey */
public Object put(Object theKey, Object theElement)
{
   BinaryTreeNode p = root,     // search pointer
                  pp = null;    // parent of p
   Comparable elementKey = (Comparable) theKey;
   // find place to insert theElement
   while (p != null)
   {// examine p.element.key
      pp = p;
      // move p to a child
      if (elementKey.compareTo(((Data) p.element).key) < 0)
         p = p.leftChild;
      else if (elementKey.compareTo(((Data) p.element).key) > 0)
              p = p.rightChild;
           else
           {// overwrite element with same key
              Object elementToReturn = ((Data) p.element).element;
              ((Data) p.element).element = theElement;
              return elementToReturn;
           }
   }

   // get a node for theElement and attach to pp
   IndexedElement q = new IndexedElement(0, theElement);
   BinaryTreeNode r = new BinaryTreeNode
                          (new Data(elementKey, q));
   if (root != null)
      // the tree is not empty
      if (elementKey.compareTo(((Data) pp.element).key) < 0)
         pp.leftChild = r;
      else
         pp.rightChild = r;
   else // insertion into empty tree
      root = r;

   // a new element has been inserted , update leftSize
   // by following search path to elementKey
   p = root;
   while (true)
      if (elementKey.compareTo(((Data) p.element).key) < 0)
      {// theElement was inserted in left subtree of p
         ((IndexedElement) ((Data) p.element).element).leftSize++;
         p = p.leftChild;
      }
      else
         if (elementKey.compareTo(((Data) p.element).key) > 0)
            // theElement was inserted in right subtree of p
            p = p.rightChild;
         else
            // reached newly inserted element
            return null;
}




An element may be removed using a two step algorithm. In the first, a search is made to verify the presence of the desired element. In the second, leftSize fields are adjusted and the element removed. The code is given below. 
/** @return matching element and remove it
  * @return null if no matching element */
public Object remove(Object theKey)
{
   Object theElement = get(theKey);
   if (theElement == null)
      // no matching element
      return null;

   // there is a matching element, remove it
   theRemove(theKey);

   return theElement;
}


The method theRemove is similar to BinarySearchTree.remove except that it knows that the delete will succeed and it also updates leftSize fields. The code is given below. 
/** remove the element with key equal to theKey
  * such an element is guaranteed to exist in the tree */
void theRemove(Object theKey)
{
   Comparable searchKey = (Comparable) theKey;

   // p will eventually point to node with key searchKey
   BinaryTreeNode p = root,    // search pointer
                  pp = null;   // parent of p

   // follow path to searchKey decrementing leftSize each time
   // we move to a left subtree
   while (!((Data) p.element).key.equals(searchKey))
   {// move to a child of p
      pp = p;
      if (searchKey.compareTo(((Data) p.element).key) < 0)
      {
         ((IndexedElement) ((Data) p.element).element).leftSize--;
         p = p.leftChild;
      }
      else
         p = p.rightChild;
   }

   // restructure tree
   // handle case when p has two children
   if (p.leftChild != null && p.rightChild != null)
   {// two children
      // convert to zero or one child case
      // find element with largest key in left subtree of p
      int ls = ((IndexedElement) ((Data) p.element).element)
                                   .leftSize - 1; // save value
      BinaryTreeNode s = p.leftChild,
                     ps = p;  // parent of s
      while (s.rightChild != null)
      {// move to larger element
         ps = s;
         s = s.rightChild;
      }

      // move largest element from s to p
      p.element = s.element;
      ((IndexedElement) ((Data) p.element).element).leftSize = ls;
      p = s;
      pp = ps;
   }

   // p has at most one child, save this child in c
   BinaryTreeNode c;
   if (p.leftChild == null)
      c = p.rightChild;
   else
      c = p.leftChild;

   // remove node p
   if (p == root) root = c;
   else
   {// is p left or right child of pp?
      if (p == pp.leftChild)
         pp.leftChild = c;
      else
         pp.rightChild = c;
   }
}


An indexed delete may be done in a similar way. The code is given below. 
/** @return element with index theIndex and remove it
  * @return null if there is no element with this index */
  public Object remove(int theIndex)
{
   Data theData = getDataElement(theIndex);
   if (theData == null)
      // no element with index theIndex
      return null;

   // there is an element with index theIndex, remove it
   theRemove(theData.key);

   return ((IndexedElement) theData.element).element;
}


Our insert and delete codes penalize successful insertions and deletions over unsuccessful ones. Since we would normally expect inserts and deletes to succeed, we may rewrite the code so as to update leftSize fields on the first pass assuming that the update is necessary. In case it turns out that the update was not to be done, a second pass is made and leftSize values reset.

The get, insert, and delete methods have complexity O(h), where h is the height of the search tree. The relevant files are IndexedBinarySearchTree.*.