blob: 55ac177f8c66a43a28754439d844eb3dd411e916 [file] [log] [blame]
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());
}
}