Java Data Structures Balanced Tree

class TreeNode {
   public TreeNode(Object o) {
      left = right = null;//  w w w  . j  av  a2 s .co m
      data = o;
   }

   public Object getData() {
      return data;
   }

   public TreeNode getLeft() {
      return left;
   }

   public TreeNode getRight() {
      return right;
   }

   public void setData(Object o) {
      data = o;
   }

   public void setLeft(TreeNode l) {
      left = l;
   }

   public void setRight(TreeNode r) {
      right = r;
   }

   public String toString() {
      return "TreeNode " + data;
   }

   TreeNode left;
   TreeNode right;
   Object data;
}

interface Comparable {
   public int compare(Object a, Object b);
}

interface Traversal {
   public void process(Object o);
}

class Tree {
   public Tree(Comparable o) {
      c = o;
   }

   public void add(Object o) {
      add(root, new TreeNode(o));
   }

   protected void add(TreeNode root, TreeNode newNode) {
      if (root == null) {
         this.root = newNode;
         return;
      }

      int val = c.compare(newNode.getData(), root.getData());

      if (val == 0) {
         root.setData(newNode.getData());
         return;
      } else if (val < 0) {
         if (root.getLeft() == null)
            root.setLeft(newNode);
         else
            add(root.getLeft(), newNode);
      } else if (val > 0) {
         if (root.getRight() == null)
            root.setRight(newNode);
         else
            add(root.getRight(), newNode);
      }
   }

   public Object search(Object o) {
      return search(root, o);
   }

   protected Object search(TreeNode root, Object o) {
      if (root == null) {
         return null;
      }

      int val = c.compare(o, root.getData());

      if (val == 0) {
         return root.getData();
      } else if (val < 0) {
         return search(root.getLeft(), o);
      } else if (val > 0) {
         return search(root.getRight(), o);
      }
      return null;
   }

   public void traverse(Traversal t) {
      traverse(INORDER, t);
   }

   public void traverse(int type, Traversal t) {
      traverse(root, type, t);
   }

   protected void traverse(TreeNode root, int type, Traversal t) {
      TreeNode tmp;

      if (type == PREORDER)
         t.process(root.getData());

      if ((tmp = root.getLeft()) != null)
         traverse(tmp, type, t);

      if (type == INORDER)
         t.process(root.getData());

      if ((tmp = root.getRight()) != null)
         traverse(tmp, type, t);
   }

   protected TreeNode root;

   protected Comparable c;

   public final static int INORDER = 1;
   public final static int PREORDER = 2;

   protected final static int RIGHT = 1;
   protected final static int LEFT = 2;
}

class BalTree extends Tree {
   public BalTree(Comparable o) {
      super(o);
   }

   public void add(Object o) {
      super.add(o);
      if (root != null) {
         root = balance(root);
      }
   }

   protected int branchCount(TreeNode root, int direction) {
      count = 0;
      TreeNode branch = null;

      if (root == null)
         return 0;

      switch (direction) {
      case RIGHT:
         branch = root.getRight();
         break;
      case LEFT:
         branch = root.getLeft();
         break;
      }

      if (branch == null)
         return 0;
      traverse(branch, INORDER, new Traversal() {
         public void process(Object o) {
            count++;
         }
      });

      return count;
   }

   protected TreeNode rotate(TreeNode root, int direction) {
      TreeNode newRoot = null;
      TreeNode orphan = null;

      switch (direction) {
      case RIGHT:
         newRoot = root.getLeft();
         root.setLeft(null);
         orphan = newRoot.getRight();
         newRoot.setRight(root);
         break;
      case LEFT:
         newRoot = root.getRight();
         root.setRight(null);
         orphan = newRoot.getLeft();
         newRoot.setLeft(root);
         break;
      }

      if (newRoot == null)
         return root;

      if (orphan != null)
         add(root, orphan);

      return newRoot;
   }

   protected TreeNode balance(TreeNode root) {
      if (root == null)
         return null;

      int left = branchCount(root, LEFT);
      int right = branchCount(root, RIGHT);
      if (left > right) {
         while (left > right + 1) {
            root = rotate(root, RIGHT);
            left = branchCount(root, LEFT);
            right = branchCount(root, RIGHT);
         }
      }

      if (right > left) {
         while (right > left + 1) {
            root = rotate(root, LEFT);
            left = branchCount(root, LEFT);
            right = branchCount(root, RIGHT);
         }
      }

      root.setLeft(balance(root.getLeft()));
      root.setRight(balance(root.getRight()));

      return root;
   }

   protected int count = 0;
}

public class Main {
   public static void main(String args[]) {
      BalTree t = new BalTree(new Comparable() {
         public int compare(Object a, Object b) {
            return ((String) a).compareTo((String) b);
         }
      });

      t.add("A");
      t.add("B");
      t.add("C");
      t.add("D");
      t.add("E");
      t.add("F");
      t.add("G");
      t.add("H");
      t.add("I");
      t.add("J");
      t.add("K");
      t.add("L");
      t.add("M");
      t.add("N");
      t.add("O");
      t.add("P");
      t.add("Q");
      t.add("R");
      t.add("S");
      t.add("T");
      t.add("U");
      t.add("V");
      t.add("W");
      t.add("X");
      t.add("Y");
      t.add("Z");

      t.traverse(Tree.PREORDER, new Traversal() {
         public void process(Object o) {
            System.out.println(o);
         }
      });
   }
}