// based on: https://algs4.cs.princeton.edu/33balanced/RedBlackBST.java.html // TODO: implement NavigableSet // RAM use in a CompressedOOPS JVM (bytes per element): // ~12 when elements are inserted in sorted order // ~13.7 when elements are inserted in random order // // (Obviously more for very small sets.) sclass HyperCompactTreeSet extends AbstractSet { // A symbol table implemented using a left-leaning red-black BST. // This is the 2-3 version. // Note: We sometimes cast the nullSentinel to A // which is technically incorrect but ok because of type erasure. // May want to fix just to be clean on the source level too. private static final boolean RED = true; private static final boolean BLACK = false; // replacement for null elements private static final new O nullSentinel; private Node root; // root of the BST int size; // size of tree set // BST helper node data type abstract sclass Node { A val; // associated data Node left() { null; } // get left subtree abstract Node setLeft(Node left); // set left subtree - return potentially replaced node Node right() { null; } // get right subtree abstract Node setRight(Node right); // set right subtree - return potentially replaced node abstract bool color(); abstract Node convertToBlack(); abstract Node convertToRed(); abstract Node invertColor(); Node convertToColor(bool color) { ret color == RED ? convertToRed() : convertToBlack(); } abstract bool isLeaf(); } // This represents a common case near a red leaf - a black node // containing a red leaf in the left slot and null in the right slot. // Combined with the direct storage of black leaf children in NonLeaf, // this is all we need to get rid of all the leaf overhead. Yay! sclass SpecialNode extends Node { A leftVal; *(A *leftVal, A *val) {} bool color() { ret BLACK; } Node convertToBlack() { this; } Node convertToRed() { ret newNode(RED, val, left(), right()); } Node invertColor() { ret convertToRed(); } Node left() { ret newLeaf(RED, leftVal); } Node setLeft(Node left) { // Can we keep the optimized representation? (Probably this // is never going to be true.) if (left != null && left.isLeaf() && left.color() == RED) ret this with leftVal = left.val; else ret newNode(BLACK, val, left, right()); } Node right() { null; } Node setRight(Node right) { if (right == null) this; ret newNode(color(), val, left(), right); } bool isLeaf() { false; } } asclass NonLeaf extends Node { // either a Node or a (sentinelled) direct user value O left, right; // color of leaf if left is a user value bool defaultLeftLeafColor() { ret BLACK; } // color of leaf if right is a user value bool defaultRightLeafColor() { ret BLACK; } Node left() { ret left == null ? null : left instanceof Node ? (Node) left : newLeaf(defaultLeftLeafColor(), (A) left); } void setLeft_noMorph(Node left) { this.left = left != null && left.isLeaf() && left.color() == defaultLeftLeafColor() ? left.val : left; } void setRight_noMorph(Node right) { this.right = right != null && right.isLeaf() && right.color() == defaultRightLeafColor() ? right.val : right; } Node setLeft(Node left) { ifndef HyperCompactTreeSet_disableSpecialNodes if (color() == BLACK && right == null && left != null && left.isLeaf() && left.color() == RED) ret new SpecialNode(left.val, val); endifndef setLeft_noMorph(left); if (left == null && right() == null) ret newLeaf(color(), val); this; } Node right() { ret right == null ? null : right instanceof Node ? (Node) right : newLeaf(defaultRightLeafColor(), (A) right); } Node setRight(Node right) { // Setting right to null may produce either a leaf or a // special node, so we just go through newNode. if (right == null && this.right != null) ret newNode(color(), val, left(), null); // New right is not null, so we compress (if possible) and store it setRight_noMorph(right); this; } bool isLeaf() { false; } } sclass BlackNode extends NonLeaf { *(A *val) {} bool color() { ret BLACK; } Node convertToBlack() { this; } Node convertToRed() { ret newNode(RED, val, left(), right()); } Node invertColor() { ret convertToRed(); } } sclass RedNode extends NonLeaf { *(A *val) {} bool color() { ret RED; } Node convertToBlack() { ret newNode(BLACK, val, left(), right()); } Node convertToRed() { this; } Node invertColor() { ret convertToBlack(); } } asclass Leaf extends Node { bool isLeaf() { true; } Node setLeft(Node left) { ret left == null ? this : newNode(color(), val, left, null); } Node setRight(Node right) { ret right == null ? this : newNode(color(), val, null, right); } } sclass BlackLeaf extends Leaf { *(A *val) {} bool color() { ret BLACK; } Node convertToBlack() { this; } Node convertToRed() { ret new RedLeaf(val); } Node invertColor() { ret convertToRed(); } } sclass RedLeaf extends Leaf { *(A *val) {} bool color() { ret RED; } Node convertToBlack() { ret new BlackLeaf(val); } Node convertToRed() { this; } Node invertColor() { ret convertToBlack(); } } *() {} *(Cl cl) { addAll(cl); } private static O deSentinel(O o) { ret o == nullSentinel ? null : o; } private static O sentinel(O o) { ret o == null ? nullSentinel : o; } // returns false on null (algorithm needs this) static bool isRed(Node x) { ret x != null && x.color() == RED; } static Node newLeaf(bool color, A val) { ret color == RED ? new RedLeaf(val) : new BlackLeaf(val); } static Node newNode(bool color, A val, Node left, Node right) { // Make leaf (always a temporary object now) if (left == null && right == null) ret newLeaf(color, val); ifndef HyperCompactTreeSet_disableSpecialNodes // Make special node if (color == BLACK && right == null && left != null && left.isLeaf() && left.color() == RED) ret new SpecialNode(left.val, val); endifndef // Make normal non-leaf NonLeaf node = color == RED ? new RedNode(val) : new BlackNode(val); node.setLeft_noMorph(left); node.setRight_noMorph(right); ret node; } public int size() { ret size; } public bool isEmpty() { ret root == null; } public bool add(A val) { val = (A) sentinel(val); int oldSize = size; root = put(root, val); root = root.convertToBlack(); ifdef CompactTreeSet_debug assertTrue(check()); endifdef ret size > oldSize; } // insert the value in the subtree rooted at h private Node put(Node h, A val) { if (h == null) { ++size; ret new RedLeaf(val); } int cmp = compare_deSentinel(val, h.val); if (cmp < 0) h = h.setLeft(put(h.left(), val)); else if (cmp > 0) h = h.setRight(put(h.right(), val)) ; else { /*h.val = val;*/ } // no overwriting // fix-up any right-leaning links if (isRed(h.right()) && !isRed(h.left())) h = rotateLeft(h); if (isRed(h.left()) && isRed(h.left().left())) h = rotateRight(h); if (isRed(h.left()) && isRed(h.right())) h = flipColors(h); ret h; } final int compare_deSentinel(A a, A b) { ret compare((A) deSentinel(a), (A) deSentinel(b)); } // override me if you wish int compare(A a, A b) { ret cmp(a, b); } public bool remove(O key) { if (!contains(key)) false; key = sentinel(key); // if both children of root are black, set root to red if (!isRed(root.left()) && !isRed(root.right())) root = root.convertToRed(); root = delete(root, (A) key); if (!isEmpty()) root = root.convertToBlack(); // assert check(); true; } // delete the key-value pair with the given key rooted at h private Node delete(Node h, A key) { // assert get(h, key) != null; if (compare_deSentinel(key, h.val) < 0) { if (!isRed(h.left()) && !isRed(h.left().left())) h = moveRedLeft(h); h = h.setLeft(delete(h.left(), key)); } else { if (isRed(h.left())) h = rotateRight(h); if (compare_deSentinel(key, h.val) == 0 && (h.right() == null)) { --size; null; } if (!isRed(h.right()) && !isRed(h.right().left())) h = moveRedRight(h); if (compare_deSentinel(key, h.val) == 0) { --size; Node x = min(h.right()); h.val = x.val; // h.val = get(h.right(), min(h.right()).val); // h.val = min(h.right()).val; h = h.setRight(deleteMin(h.right())); } else h = h.setRight(delete(h.right(), key)); } return balance(h); } // make a left-leaning link lean to the right private Node rotateRight(Node h) { // assert (h != null) && isRed(h.left()); Node x = h.left(); h = h.setLeft(x.right()); x = x.setRight(h); x = x.convertToColor(x.right().color()); x = x.setRight(x.right().convertToRed()); ret x; } // make a right-leaning link lean to the left private Node rotateLeft(Node h) { // assert (h != null) && isRed(h.right()); Node x = h.right(); h = h.setRight(x.left()); x = x.setLeft(h); x = x.convertToColor(x.left().color()); x = x.setLeft(x.left().convertToRed()); ret x; } // flip the colors of a node and its two children private Node flipColors(Node h) { // h must have opposite color of its two children // assert (h != null) && (h.left() != null) && (h.right() != null); // assert (!isRed(h) && isRed(h.left()) && isRed(h.right())) // || (isRed(h) && !isRed(h.left()) && !isRed(h.right())); h = h.setLeft(h.left().invertColor()); h = h.setRight(h.right().invertColor()); ret h.invertColor(); } // Assuming that h is red and both h.left() and h.left().left() // are black, make h.left() or one of its children red. private Node moveRedLeft(Node h) { // assert (h != null); // assert isRed(h) && !isRed(h.left()) && !isRed(h.left().left()); h = flipColors(h); if (isRed(h.right().left())) { h = h.setRight(rotateRight(h.right())); h = rotateLeft(h); h = flipColors(h); } ret 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. private Node moveRedRight(Node h) { // assert (h != null); // assert isRed(h) && !isRed(h.right()) && !isRed(h.right().left()); h = flipColors(h); if (isRed(h.left().left())) { h = rotateRight(h); h = flipColors(h); } ret h; } // restore red-black tree invariant private Node balance(Node h) { // assert (h != null); 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())) h = flipColors(h); ret h; } /** * Returns the height of the BST (for debugging). * @return the height of the BST (a 1-node tree has height 0) */ public int height() { ret height(root); } private int height(Node x) { if (x == null) return -1; return 1 + Math.max(height(x.left()), height(x.right())); } public bool contains(O val) { ret find(root, (A) sentinel(val)) != null; } public A find(A probeVal) { probeVal = (A) sentinel(probeVal); Node n = find(root, probeVal); ret n == null ? null : n.val; } // value associated with the given key in subtree rooted at x; null if no such key private A get(Node x, A key) { x = find(x, key); ret x == null ? null : x.val; } Node find(Node x, A key) { while (x != null) { int cmp = compare_deSentinel(key, x.val); if (cmp < 0) x = x.left(); else if (cmp > 0) x = x.right(); else ret x; } null; } private boolean check() { if (!is23()) println("Not a 2-3 tree"); if (!isBalanced()) println("Not balanced"); return is23() && isBalanced(); } // Does the tree have no red right links, and at most one (left) // red links in a row on any path? private boolean is23() { return is23(root); } private boolean is23(Node x) { if (x == null) true; if (isRed(x.right())) false; if (x != root && isRed(x) && isRed(x.left())) false; ret is23(x.left()) && is23(x.right()); } // do all paths from root to leaf have same number of black edges? private bool isBalanced() { 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(); } ret isBalanced(root, black); } // does every path from the root to a leaf have the given number of black links? private boolean isBalanced(Node x, int black) { if (x == null) return black == 0; if (!isRed(x)) black--; return isBalanced(x.left(), black) && isBalanced(x.right(), black); } public void clear() { root = null; size = 0; } // the smallest key in subtree rooted at x; null if no such key private Node min(Node x) { // assert x != null; while (x.left() != null) x = x.left(); ret x; } private Node deleteMin(Node h) { if (h.left() == null) return null; if (!isRed(h.left()) && !isRed(h.left().left())) h = moveRedLeft(h); h = h.setLeft(deleteMin(h.left())); ret balance(h); } public Iterator iterator() { ret new MyIterator; } class MyIterator extends ItIt { new L> path; *() { fetch(root); } void fetch(Node node) { while (node != null) { path.add(node); node = node.left(); } } public bool hasNext() { ret !path.isEmpty(); } public A next() { if (path.isEmpty()) fail("no more elements"); Node node = popLast(path); // last node is always a leaf, so left is null // so proceed to fetch right branch fetch(node.right()); ret (A) deSentinel(node.val); } } // Returns the smallest key in the symbol table greater than or equal to {@code key}. public A ceiling(A key) { key = (A) sentinel(key); Node x = ceiling(root, key); ret x == null ? null : x.val; } // the smallest key in the subtree rooted at x greater than or equal to the given key Node ceiling(Node x, A key) { if (x == null) null; int cmp = compare_deSentinel(key, x.val); if (cmp == 0) ret x; if (cmp > 0) ret ceiling(x.right(), key); Node t = ceiling(x.left(), key); if (t != null) ret t; else ret x; } public A floor(A key) { key = (A) sentinel(key); Node x = floor(root, key); ret x == null ? null : x.val; } // the largest key in the subtree rooted at x less than or equal to the given key Node floor(Node x, A key) { if (x == null) null; int cmp = compare_deSentinel(key, x.val); if (cmp == 0) ret x; if (cmp < 0) ret floor(x.left(), key); Node t = floor(x.right(), key); if (t != null) ret t; else ret x; } void testInternalStructure() { ifndef HyperCompactTreeSet_disableSpecialNodes // one leaf object (root) is allowed - could even optimize that assertTrue(countLeafObjects() <= 1); endifndef } // count leaf objects we didn't optimize away int countLeafObjects(Node node default root) { if (node instanceof Leaf) ret 1; if (node cast NonLeaf) ret countLeafObjects(optCast Node(node.left)) + countLeafObjects(optCast Node(node.right)); ret 0; } Cl unoptimizedNodes() { new L out; findUnoptimizedNodes(out); ret out; } void findUnoptimizedNodes(Node node default root, L out) { if (node == null) ret; if (node cast NonLeaf) { if (isUnoptimizedNode(node)) out.add(node); findUnoptimizedNodes(optCast Node(node.left), out); findUnoptimizedNodes(optCast Node(node.right), out); } } bool isUnoptimizedNode(Node node) { if (node cast NonLeaf) ret node.left instanceof Leaf || node.right instanceof Leaf; false; } // Compact me (minimize space use). Returns a compacted copy. // This only has an effect when elements were inserted in non-sorted // or and/or elements were removed. selfType compact() { ret new HyperCompactTreeSet(this); } }