Java Data Structures Red Black Tree

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

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

class TreeNode {
   public TreeNode(Object o) {
      data = o;/*from  w w w .  j  a v a2s  .  c  o m*/
      color = RED;
      left = right = null;
   }

   public Object getData() {
      return data;
   }

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

   public TreeNode getLeft() {
      return left;
   }

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

   public TreeNode getRight() {
      return right;
   }

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

   public boolean getColor() {
      return color;
   }

   public void setColor(boolean c) {
      color = c;
   }

   public void flip() {
      color = !color;
   }

   public boolean hasRedChild() {
      if (left != null && left.color == RED)
         return true;

      if (right != null && right.color == RED)
         return true;

      return false;
   }

   public boolean is2Way() {
      if (color == RED)
         return false;

      return !hasRedChild();
   }

   public boolean is3Way() {
      if (color == RED)
         return false;

      if (is2Way() || is4Way())
         return false;

      return true;
   }

   public boolean is4Way() {
      if (color == RED)
         return false;

      if (left == null || right == null)
         return false;

      if (left.color == RED && right.color == RED)
         return true;

      return false;
   }

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

   private TreeNode left;
   private TreeNode right;
   private Object data;
   private boolean color;

   public static final boolean RED = true;
   public static final boolean BLACK = false;
}

class RBTree {
   public RBTree(Comparable c) {
      this.c = c;
      root = null;
   }

   public void add(Object o) {
      root = add(root, new TreeNode(o));
      root.setColor(TreeNode.BLACK);
   }

   protected TreeNode add(TreeNode root, TreeNode newNode) {
      if (root == null)
         return newNode;

      if (root.is4Way())
         split(root);

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

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

      root = balance(root);
      return root;
   }

   protected TreeNode balance(TreeNode node) {
      if (node.hasRedChild() == false)
         return node;

      TreeNode child = node.getLeft();

      if (child != null) {
         if (child.hasRedChild() == true)
            node = rotate(node, RIGHT);
      }

      child = node.getRight();

      if (child != null) {
         if (child.hasRedChild() == true)
            node = rotate(node, LEFT);
      }

      return node;
   }

   public void split(TreeNode node) {
      node.flip();
      node.getRight().flip();
      node.getLeft().flip();
   }

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

      tmp = root.getColor();
      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;

      root.setColor(newRoot.getColor());
      newRoot.setColor(tmp);
      if (orphan != null)
         add(root, orphan);

      return newRoot;
   }

   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 TreeNode lastBlack;
   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;
}

public class Main {
   public static void main(String args[]) {
      RBTree t = new RBTree(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(RBTree.PREORDER, new Traversal() {
         public void process(Object o) {
            System.out.println(o);
         }
      });
   }
}