Sedgewick-RedBlackBST.java

public class RedBlackBST<Key extends Comparable<Key>, Value> 
{
   private static final int BST = 0;
   private static final int TD234 = 1;
   private static final int BU23 = 2;
   private static final boolean RED   = true;
   private static final boolean BLACK = false;

   private Node root;            // root of the BST
   private int k;                // ordinal for drawing
   private final int species;    // species kind of tree for insert
   private int heightBLACK;      // black height of tree 

   RedBlackBST(int species)
   {  this.species = species;  }

   private class Node
   {
      Key key;                  // key
      Value value;              // associated data
      Node left, right;         // left and right subtrees
      boolean color;            // color of parent link
      private int N;            // number of nodes in tree rooted here
      private int height;       // height of tree rooted here
      private double xc, yc;    // for drawing

      Node(Key key, Value value)
      {
         this.key   = key;
         this.value = value;
         this.color = RED;
         this.N = 1;
         this.height = 1;
      }
   }

   public int size()
   {  return size(root);  }

   private int size(Node x)
   { 
      if (x == null) return 0;
      else           return x.N;
   }

   public int rootRank()
   { 
       if (root == null) return 0;
       else              return size(root.left);
   }

   public int height()
   {  return height(root);  }

   public int heightB()
   {  return heightBLACK;  }

   private int height(Node x)
   { 
      if (x == null) return 0;
      else           return x.height;
   }

   public boolean contains(Key key)
   {  return (get(key) != null);  }

   public Value get(Key key)
   {  return get(root, key);  }

   private Value get(Node x, Key key)
   {
      if (x == null)        return null;
      if (eq  (key, x.key)) return x.value;
      if (less(key, x.key)) return get(x.left,  key);
      else                  return get(x.right, key);
   }

   public Key min()
   {  
      if (root == null) return null;
      else              return min(root);
   }

   private Key min(Node x)
   {
      if (x.left == null) return x.key;
      else                return min(x.left);
   }

   public Key max()
   {  
      if (root == null) return null;
      else              return max(root);
   }

   private Key max(Node x)
   {
      if (x.right == null) return x.key;
      else                 return max(x.right);
   }

   public void put(Key key, Value value)
   {
      root = insert(root, key, value);
      if (isRed(root)) heightBLACK++;
      root.color = BLACK;
   }

   private Node insert(Node h, Key key, Value value)
   { 
      if (h == null) 
         return new Node(key, value);

      if (species == TD234) 
         if (isRed(h.left) && isRed(h.right))
	     colorFlip(h);

      if (eq(key, h.key))
         h.value = value;
      else if (less(key, h.key)) 
         h.left = insert(h.left, key, value); 
      else 
         h.right = insert(h.right, key, value); 

      if (species == BST) return setN(h);

      if (isRed(h.right))
         h = rotateLeft(h);

      if (isRed(h.left) && isRed(h.left.left))
         h = rotateRight(h);

      if (species == BU23)
         if (isRed(h.left) && isRed(h.right))
	     colorFlip(h);

      return setN(h);
   }

   public void deleteMin()
   {
      root = deleteMin(root);
      root.color = BLACK;
   }

   private Node deleteMin(Node h)
   { 
      if (h.left == null)
         return null;

      if (!isRed(h.left) && !isRed(h.left.left))
         h = moveRedLeft(h);

      h.left = deleteMin(h.left);

      return fixUp(h);
   }

   public void deleteMax()
   {
      root = deleteMax(root);
      root.color = BLACK;
   }

   private Node deleteMax(Node h)
   { 
       //      if (h.right == null)
	  //      {  
	  //         if (h.left != null)
	  //            h.left.color = BLACK;
	  //         return h.left;
	  //      }

      if (isRed(h.left))
         h = rotateRight(h);

      if (h.right == null)
         return null;

      if (!isRed(h.right) && !isRed(h.right.left))
         h = moveRedRight(h);

      h.right = deleteMax(h.right);

      return fixUp(h);
   }

   public void delete(Key key)
   { 
      root = delete(root, key);
      root.color = BLACK;
   }

   private Node delete(Node h, Key key)
   { 
      if (less(key, h.key)) 
      {
         if (!isRed(h.left) && !isRed(h.left.left))
            h = moveRedLeft(h);
         h.left =  delete(h.left, key);
      }
      else 
      {
         if (isRed(h.left)) 
            h = rotateRight(h);
         if (eq(key, h.key) && (h.right == null))
            return null;
         if (!isRed(h.right) && !isRed(h.right.left))
            h = moveRedRight(h);
         if (eq(key, h.key))
         {
            h.value = get(h.right, min(h.right));
            h.key = min(h.right);
	    h.right = deleteMin(h.right);
         }
         else h.right = delete(h.right, key);
      }
 
      return fixUp(h);
   }

// Helper methods

    private boolean less(Key a, Key b) { return a.compareTo(b) <  0; }
    private boolean eq  (Key a, Key b) { return a.compareTo(b) == 0; }

    private boolean isRed(Node x)
    {
        if (x == null) return false;
        return (x.color == RED);
    }

   private void colorFlip(Node h)
   {  
      h.color        = !h.color;
      h.left.color   = !h.left.color;
      h.right.color  = !h.right.color;
   }

   private Node rotateLeft(Node h)
   {  // Make a right-leaning 3-node lean to the left.
      Node x = h.right;
      h.right = x.left;
      x.left = setN(h);
      x.color      = x.left.color;                   
      x.left.color = RED;                     
      return setN(x);
   }

   private Node rotateRight(Node h)
   {  // Make a left-leaning 3-node lean to the right.
      Node x = h.left;
      h.left = x.right;
      x.right = setN(h);
      x.color       = x.right.color;                   
      x.right.color = RED;                     
      return setN(x);
    }

    private Node moveRedLeft(Node h)
    {  // Assuming that h is red and both h.left and h.left.left
       // are black, make h.left or one of its children red.
       colorFlip(h);
       if (isRed(h.right.left))
       { 
          h.right = rotateRight(h.right);
          h = rotateLeft(h);
          colorFlip(h);
       }
      return h;
    }

    private Node moveRedRight(Node h)
    {  // Assuming that h is red and both h.right and h.right.left
       // are black, make h.right or one of its children red.
       colorFlip(h);
       if (isRed(h.left.left))
       { 
          h = rotateRight(h);
	  colorFlip(h);
       }
       return h;
    }

   private Node fixUp(Node h)
   {
      if (isRed(h.right))
         h = rotateLeft(h);

      if (isRed(h.left) && isRed(h.left.left))
         h = rotateRight(h);

      if (isRed(h.left) && isRed(h.right))
         colorFlip(h);

      return setN(h);
   }
      
   private Node setN(Node h)
   {
      h.N = size(h.left) + size(h.right) + 1;
      if (height(h.left) > height(h.right)) h.height = height(h.left) + 1;
      else                                  h.height = height(h.right) + 1;
      return h;
   }
      
   public String toString()
   {  
      if (root == null) return "";
      else              return heightB() + " " + toString(root);  
   }

   public String toString(Node x)
   {  
      String s = "(";
      if (x.left == null) s += "("; else s += toString(x.left);
      if (isRed(x)) s += "*";
      if (x.right == null) s += ")"; else s += toString(x.right);
      return s + ")";
   }

// Methods for tree drawing

    public void draw(double y, double lineWidth, double nodeSize)
    { 
      k = 0; 
      setcoords(root, y);
      StdDraw.setPenColor(StdDraw.BLACK);
      StdDraw.setPenRadius(lineWidth);
      drawlines(root); 
      StdDraw.setPenColor(StdDraw.WHITE);
      drawnodes(root, nodeSize); 
    }

    public void setcoords(Node x, double d)
    {
        if (x == null) return;
        setcoords(x.left, d-.04);
        x.xc = (0.5 + k++)/size(); x.yc = d - .04;
        setcoords(x.right, d-.04);
    }

    public void drawlines(Node x)
    {
        if (x == null) return;
        drawlines(x.left);
        if (x.left != null)
        {
	    if (x.left.color == RED) StdDraw.setPenColor(StdDraw.RED);
	    else StdDraw.setPenColor(StdDraw.BLACK);
            StdDraw.line(x.xc, x.yc, x.left.xc, x.left.yc);
        }
        if (x.right != null)
        {
	    if (x.right.color == RED) StdDraw.setPenColor(StdDraw.RED);
	    else StdDraw.setPenColor(StdDraw.BLACK);
            StdDraw.line(x.xc, x.yc, x.right.xc, x.right.yc);
        }
        drawlines(x.right);
    }

    public void drawnodes(Node x, double nodeSize)
    {
        if (x == null) return;
        drawnodes(x.left, nodeSize);
	StdDraw.filledCircle(x.xc, x.yc, nodeSize);
        drawnodes(x.right, nodeSize);
    }

    public void mark(Key key)
    { 
      StdDraw.setPenColor(StdDraw.BLACK);
      marknodes(key, root); 
    }

    public void marknodes(Key key, Node x)
    {
        if (x == null) return;
        marknodes(key, x.left);
	if (eq(key, x.key))
           StdDraw.filledCircle(x.xc, x.yc, .004);
        marknodes(key, x.right);
    }

    public int ipl()
    {  return ipl(root);  }

    public int ipl(Node x)
    {
	if (x == null) return 0;
	return size(x) - 1 + ipl(x.left) + ipl(x.right);
    }
       
    public int sizeRed()
    {  return sizeRed(root);  }

    public int sizeRed(Node x)
    {
	if (x == null) return 0;
        if (isRed(x)) return 1 + sizeRed(x.left) + sizeRed(x.right);
        else          return sizeRed(x.left) + sizeRed(x.right);
    }
       
// Integrity checks

   public boolean check() 
   {  // Is this tree a red-black tree?
      return isBST() && is234() && isBalanced();
   }

   private boolean isBST() 
   {  // Is this tree a BST?
      return isBST(root, min(), max());
   }

   private boolean isBST(Node x, Key min, Key max)
   {  // Are all the values in the BST rooted at x between min and max,
      // and does the same property hold for both subtrees?
      if (x == null) return true;
      if (less(x.key, min) || less(max, x.key)) return false;
      return isBST(x.left, min, x.key) && isBST(x.right, x.key, max);
   } 

   private boolean is234() { return is234(root); }
   private boolean is234(Node x)
   {  // Does the tree have no red right links, and at most two (left)
      // red links in a row on any path?
      if (x == null) return true;
      if (isRed(x.right)) return false;
      if (isRed(x))
        if (isRed(x.left))
          if (isRed(x.left.left)) return false;
      return is234(x.left) && is234(x.right);
   } 

   private boolean isBalanced()
   { // Do all paths from root to leaf have same number of black edges?
      int black = 0;     // number of black links on path from root to min
      Node x = root;
      while (x != null)
      {
         if (!isRed(x)) black++;
            x = x.left;
      }
      return isBalanced(root, black);
   }

   private boolean isBalanced(Node x, int black)
   { // Does every path from the root to a leaf have the given number 
     // of black links?
      if      (x == null && black == 0) return true;
      else if (x == null && black != 0) return false;
      if (!isRed(x)) black--;
      return isBalanced(x.left, black) && isBalanced(x.right, black);
   } 


  public static void main(String[] args)
  {
      StdDraw.setPenRadius(.0025);
      int species = Integer.parseInt(args[0]);
      RedBlackBST<Integer, Integer> st;
      st = new RedBlackBST<Integer, Integer>(species);
      int[] a = { 3, 1, 4, 2, 5, 9, 6, 8, 7 };
      for (int i = 0; i < a.length; i++)
	  st.put(a[i], i);
      StdOut.println(st);
      StdDraw.clear(StdDraw.LIGHT_GRAY);
      st.draw(.95, .0025, .008);
      StdOut.println(st.min() + " " + st.max() + " " + st.check());
      StdOut.println(st.ipl());
      StdOut.println(st.heightB());
  }

}