public class RedBlackBST, 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 st; st = new RedBlackBST(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()); } }