Simple Binary Tree Class

/*
 *  Title: class BinaryTree
 *  Description: Simple binary tree class to demonstrate the basic principles
 *  of implementing a tree data structure. This should not be taken as a production
 *  quality class (see the text book instead).
 *  Copyright:   Copyright (c) 2002
 *  Company:     Dept. of Computer Science, University College London
 *  @author Graham Roberts
 *  @version 1.0 01-Mar-02
 */
import java.util.Stack;
public class BinaryTree
{
  /**
   * A tree is a hierarchical structure of TreeNode objects. root references
   * the first node on the tree.
   */
  private TreeNode root;
  /**
   *  Helper class used to implement tree nodes. As this is a private helper
   *  class it is acceptable to have public instance variables. Instances of
   *  this class are never made available to client code of the tree.
   */
  private static class TreeNode
  {
    /**
     *  Data object reference.
     */
    public Comparable val;
    /**
     *  Left and right child nodes.
     */
    public TreeNode left,right;
    /**
     *  Constructor for TreeNode.
     *
     *@param  val    data object reference
     *@param  left   left child node reference or null
     *@param  right  right child node reference or null
     */
    public TreeNode(Comparable val, TreeNode left, TreeNode right)
    {
      this.val = val;
      this.left = left;
      this.right = right;
    }
    /**
     * Insert an object into the tree.
     *
     * @param obj object to insert into tree.
     */
    public void insert(Comparable obj)
    {
      if (val.compareTo(obj) < 0)
      {
        if (right != null)
        {
          right.insert(obj) ;
        }
        else
        {
          right = new TreeNode(obj,null,null) ;
        }
      }
      else
      {
        if (left != null)
        {
          left.insert(obj) ;
        }
        else
        {
          left = new TreeNode(obj,null,null) ;
        }
      }
    }
    /**
     * Find an object in the tree. Objects are compared using the compareTo method, so
     * must conform to type Comparable.
     * Two objects are equal if they represent the same value.
     *
     * @param obj Object representing value to find in tree.
     * @return  reference to matching node or null.
     */
    public TreeNode find(Comparable obj)
    {
      int temp = val.compareTo(obj) ;
      if (temp == 0)
      {
        return this ;
      }
      if (temp < 0)
      {
        return (right == null) ? null : right.find(obj) ;
      }
      return (left == null) ? null : left.find(obj) ;
    }
    /**
     * Remove the node referencing an object representing the same value as the argument object.
     * This recursive method essentially restructures the tree as necessary and returns a
     * reference to the new root. The algorithm is straightforward apart from the case
     * where the node to be removed has two children. In that case the left-most leaf node
     * of the right child is moved up the tree to replace the removed node. Hand work some
     * examples to see how this works.
     *
     * @param obj Object representing value to remove from tree.
     * @param t Root node of the sub-tree currently being examined (possibly null).
     * @return reference to the (possibly new) root node of the sub-tree being examined or
     * null if no node.
     */
    private TreeNode remove(Comparable obj, TreeNode t)
    {
      if (t == null)
      {
        return t;
      }
      if (obj.compareTo(t.val) < 0)
      {
        t.left = remove(obj,t.left);
      }
      else
      if (obj.compareTo(t.val) > 0 )
      {
        t.right = remove(obj, t.right);
      }
      else
      if (t.left != null && t.right != null)
      {
        t.val = findMin(t.right).val;
        t.right = remove(t.val,t.right);
      }
      else
      {
        t = (t.left != null) ? t.left : t.right;
      }
      return t;
    }
    /**
     * Helper method to find the left most leaf node in a sub-tree.
     *
     * @param t TreeNode to be examined.
     * @return reference to left most leaf node or null.
     */
    private TreeNode findMin(TreeNode t)
    {
      if(t == null)
      {
        return null;
      }
      else
      if(t.left == null)
      {
        return t;
      }
      return findMin(t.left);
    }
  }
  /**
   * Construct an empty tree.
   */
  public BinaryTree()
  {
    root = null ;
  }
  /**
   * Store an object in the tree. The object must conform to type Comparable
   * in order to be inserted in the correct location. Multiple objects representing the
   * same value can be added.
   *
   * @param obj reference to Comparable object to add.
   */
  public void add(Comparable obj)
  {
    if (root == null)
    {
      root = new TreeNode(obj,null,null) ;
    }
    else
    {
      root.insert(obj) ;
    }
  }
  /**
   * Determine whether the tree contains an object with the same value as the
   * argument.
   *
   * @param obj reference to Comparable object whose value will be searched for.
   * @return true if the value is found.
   */
  public boolean contains(Comparable obj)
  {
    if (root == null)
    {
      return false ;
    }
    else
    {
      return (root.find(obj) != null) ;
    }
  }
  /**
   * Remove an object with a matching value from the tree. If multiple
   * matches are possible, only the first matching object is removed.
   *
   * @param obj Remove an object with a matching value from the tree.
   */
  public void remove(Comparable obj)
  {
    if (root != null)
    {
      root = root.remove(obj,root) ;
    }
  }
  /**
   * Simple pre-order iterator class. An iterator object will sequence through
   * the tree contents in ascending order.
   * A stack is used to keep track of where the iteration has reached in the tree.
   * Note that if new items are added or removed during an iteration, there is no
   * guarantee that the iteration will continue correctly.
   */
  public class Iterator
  {
    private Stack nodes = new Stack() ;
    /**
     * Construct a new iterator for the current tree object.
     */
    public Iterator()
    {
      pushLeft(root) ;
    }
    /**
     * Get next obnject in sequence.
     *
     * @return next object in sequence or null if the end of the sequence has
     * been reached.
     */
    public Comparable next()
    {
      if (nodes.isEmpty())
      {
        return null ;
      }
      TreeNode node = (TreeNode)nodes.pop() ;
      pushLeft(node.right) ;
      return node.val ;
    }
    /**
     * Determine if there is another object in the iteration sequence.
     *
     * @return true if another object is available in the sequence.
     */
    public boolean hasNext()
    {
      return !nodes.isEmpty() ;
    }
    /**
     * Helper method used to push node objects onto the stack to keep
     * track of the iteration.
     */
    private void pushLeft(TreeNode node)
    {
      while (node != null)
      {
        nodes.push(node) ;
        node = node.left ;
      }
    }
  }
  /**
   * Return a new tree iterator object.
   *
   * @return new iterator object.
   */
  public Iterator iterator()
  {
    return new Iterator() ;
  }
}
 

JUnit test class for class BinaryTree

/*
 * Title: class BinaryTreeTest
 * Description: JUnit test class for BinaryTree class
 * Copyright:   Copyright (c) 2002
 * Company:     Dept. of Computer Science, University College London
 * @author Graham Roberts
 * @version 1.0 01-Mar-02
 */
import junit.framework.* ;
public class BinaryTreeTest extends TestCase
{
  private BinaryTree empty ;
  private BinaryTree one ;
  private BinaryTree several ;
  public BinaryTreeTest(String s)
  {
    super(s) ;
  }
  public void setUp()
  {
    empty = new BinaryTree() ;
    one = new BinaryTree() ;
    one.add(new Integer(0)) ;
    several = new BinaryTree() ;
    several.add(new Integer(5)) ;
    several.add(new Integer(2)) ;
    several.add(new Integer(1)) ;
    several.add(new Integer(9)) ;
    several.add(new Integer(8)) ;
    several.add(new Integer(10)) ;
  }
  public void testContains()
  {
    assertTrue("Check empty is empty",!empty.contains(new Integer(0))) ;
    assertTrue("Check one contains 0",one.contains(new Integer(0))) ;
    checkTree(several,new int[]{1,2,5,8,9,10}) ;
  }
  public void testDoesNotContain()
  {
    assertTrue("Check -2 not in several",!several.contains(new Integer(-2))) ;
    assertTrue("Check 3 not in several",!several.contains(new Integer(3))) ;
    assertTrue("Check 15 not in several",!several.contains(new Integer(15))) ;
  }
  public void testRemoveFromOne()
  {
    one.remove(new Integer(0)) ;
    assertTrue("Check one is empty",!one.contains(new Integer(0))) ;
  }
  public void testRemoveLeaf()
  {
    several.remove(new Integer(1)) ;
    assertTrue("Check 1 removed from several",!several.contains(new Integer(1))) ;
    several.remove(new Integer(8)) ;
    assertTrue("Check 8 removed from several",!several.contains(new Integer(8))) ;
    several.remove(new Integer(10)) ;
    assertTrue("Check 10 removed from several",!several.contains(new Integer(10))) ;
  }
  public void testRemoveRootElement()
  {
    several.remove(new Integer(5)) ;
    assertTrue("Check 5 at root removed from several",!several.contains(new Integer(5))) ;
    checkTree(several,new int[]{1,2,8,9,10}) ;
    several.remove(new Integer(8)) ;
    assertTrue("Check 8 at root removed from several",!several.contains(new Integer(8))) ;
    checkTree(several,new int[]{1,2,9,10}) ;
    several.remove(new Integer(9)) ;
    assertTrue("Check 9 at root removed from several",!several.contains(new Integer(9))) ;
    checkTree(several,new int[]{1,2,10}) ;
    several.remove(new Integer(10)) ;
    assertTrue("Check 10 at root removed from several",!several.contains(new Integer(10))) ;
    checkTree(several,new int[]{1,2}) ;
    several.remove(new Integer(2)) ;
    assertTrue("Check 2 at root removed from several",!several.contains(new Integer(2))) ;
    checkTree(several,new int[]{1}) ;
  }
  public void testDuplicates()
  {
    empty.add(new Integer(1)) ;
    empty.add(new Integer(1)) ;
    empty.add(new Integer(1)) ;
    int[] content = new int[] {1,1,1} ;
    checkIterator(empty,content);
    assertTrue("contains 1",empty.contains(new Integer(1))) ;
    empty.remove(new Integer(1)) ;
    assertTrue("still contains 1",empty.contains(new Integer(1))) ;
    empty.remove(new Integer(1)) ;
    assertTrue("still contains 1",empty.contains(new Integer(1))) ;
    empty.remove(new Integer(1)) ;
    assertTrue("empty empty",!empty.contains(new Integer(1))) ;
  }
  private void checkTree(BinaryTree tree, int[] elements)
  {
    for (int i = 0 ; i < elements.length ; i++)
    {
      assertTrue("Check " + elements[i] + " still in tree",tree.contains(new Integer(elements[i]))) ;
    }
  }
  public void testIterator()
  {
    BinaryTree.Iterator iterator = empty.iterator() ;
    assertTrue("Nothing to iterate in empty",!iterator.hasNext()) ;
    int[] content = new int[] {0} ;
    checkIterator(one,content);
    content = new int[] {1,2,5,8,9,10} ;
    checkIterator(several,content);
  }
  private void checkIterator(BinaryTree tree, int[] content)
  {
    BinaryTree.Iterator iterator = tree.iterator() ;
    for (int i = 0 ; i < content.length ; i++)
    {
      assertEquals("Check " + content[i] + " in several tree",new Integer(content[i]),iterator.next()) ;
    }
  }
  public static Test suite()
  {
    return new TestSuite(BinaryTreeTest.class) ;
  }
  public static void main (String[] args)
  {
    junit.textui.TestRunner.run(suite()) ;
  }
}