com.ajjpj.afoundation.collection.immutable.ARedBlackTreeMap Maven / Gradle / Ivy
The newest version!
package com.ajjpj.afoundation.collection.immutable;
import com.ajjpj.afoundation.collection.AEquality;
import java.util.ArrayDeque;
import java.util.Comparator;
import java.util.Iterator;
import java.util.NoSuchElementException;
/**
* @author arno
*/
public class ARedBlackTreeMap extends AbstractAMap implements ASortedMap {
final Tree root;
private final Comparator comparator;
public static ARedBlackTreeMap empty (Comparator comparator) {
return new ARedBlackTreeMap<> (null, comparator);
}
private ARedBlackTreeMap (Tree root, Comparator comparator) {
this.root = root;
this.comparator = comparator;
}
@Override public AEquality keyEquality () {
return new AEquality.ComparatorBased (comparator);
}
@Override public AMap clear () {
return empty (comparator);
}
@Override public int size () {
return root == null ? 0 : root.count;
}
@Override public boolean containsKey (K key) {
return lookup (root, key, comparator) != null;
}
@Override public AOption get (K key) {
final Tree raw = lookup (root, key, comparator);
if (raw == null) {
return AOption.none ();
}
return AOption.some (raw.value);
}
@Override public ASet keys () {
return ARedBlackTreeSet.create (this);
}
@Override public ARedBlackTreeMap updated (K key, V value) {
return new ARedBlackTreeMap<> (blacken (upd (root, key, value, comparator)), comparator);
}
@Override public ARedBlackTreeMap removed (K key) {
return new ARedBlackTreeMap<> (blacken (del (root, key, comparator)), comparator);
}
@Override public Iterator> iterator () {
return new TreeIterator> () {
@Override AMapEntry nextResult (Tree tree) {
return tree;
}
};
}
@Override public AOption> first () {
if (root == null) return AOption.none ();
Tree cur = root;
while (cur.left != null) cur = cur.left;
return AOption.>some (cur);
}
@Override public AOption> last () {
if (root == null) return AOption.none ();
Tree cur = root;
while (cur.right != null) cur = cur.right;
return AOption.>some (cur);
}
@Override public AOption> firstGreaterThan (K key) {
if (root == null) return AOption.none ();
Tree cur = root;
Tree candidate = null;
while (true) {
final int cmp = comparator.compare (cur.key, key);
if (cmp <= 0) {
// this node is smaller than the key --> go right
if (cur.right == null) return AOption.>fromNullable (candidate);
cur = cur.right;
}
else {
// this node is greater than the key --> go left
if (cur.left == null) return AOption.>some (cur);
candidate = cur;
cur = cur.left;
}
}
}
@Override public AOption> firstGreaterOrEquals (K key) {
if (root == null) return AOption.none ();
Tree cur = root;
Tree candidate = null;
while (true) {
final int cmp = comparator.compare (cur.key, key);
if (cmp == 0) return AOption.>some (cur);
if (cmp < 0) {
// this node is smaller than the key --> go right
if (cur.right == null) return AOption.>fromNullable (candidate);
cur = cur.right;
}
else {
// this node is greater than the key --> go left
if (cur.left == null) return AOption.>some (cur);
candidate = cur;
cur = cur.left;
}
}
}
@Override public AOption> lastSmallerThan (K key) {
if (root == null) return AOption.none ();
Tree cur = root;
Tree candidate = null;
while (true) {
final int cmp = comparator.compare (cur.key, key);
if (cmp >= 0) {
// this node is greater than the key --> go left
if (cur.left == null) return AOption.>fromNullable (candidate);
cur = cur.left;
}
else {
// this node is smaller than the key --> go right
if (cur.right == null) return AOption.>some (cur);
candidate = cur;
cur = cur.right;
}
}
}
@Override public AOption> lastSmallerOrEquals (K key) {
if (root == null) return AOption.none ();
Tree cur = root;
Tree candidate = null;
while (true) {
final int cmp = comparator.compare (cur.key, key);
if (cmp == 0) return AOption.>some (cur);
if (cmp > 0) {
// this node is greater than the key --> go left
if (cur.left == null) return AOption.>fromNullable (candidate);
cur = cur.left;
}
else {
// this node is smaller than the key --> go right
if (cur.right == null) return AOption.>some (cur);
candidate = cur;
cur = cur.right;
}
}
}
private Iterator> iterator (final K from, final K to, boolean fromInclusive, boolean toInclusive) {
if (root == null) {
return new TreeIterator> (null, null) {
@Override AMapEntry nextResult (Tree tree) {
return null;
}
};
}
// this stack contains all nodes for which the left side was visited, while they themselves were not, and neither was their right side
final ArrayDeque> pathStack = new ArrayDeque<> ();
Tree cur = root;
while (true) {
final int cmp = (from == null) ? 1 : comparator.compare (cur.key, from);
if (cmp == 0 && fromInclusive) {
pathStack.push (cur);
break;
}
if (cmp <= 0) { // cmp == 0 can only happen if !fromInclusive
// this node is smaller than the key --> go right
if (cur.right == null) {
break;
}
// this node is not in range --> do not push
cur = cur.right;
}
else {
pathStack.push (cur);
// this node is greater than the key --> go left
if (cur.left == null) {
break;
}
cur = cur.left;
}
}
if (to == null) {
return new TreeIterator> (pathStack, pathStack.isEmpty () ? null : pathStack.pop ()) {
@Override AMapEntry nextResult (Tree tree) {
return tree;
}
};
}
Tree first = null;
if (! pathStack.isEmpty ()) {
first = pathStack.pop ();
final int cmp = comparator.compare (first.key, to);
if (cmp > 0 || (cmp == 0 && !toInclusive)) {
first = null;
}
}
if (toInclusive) {
return new TreeIterator> (pathStack, first) {
@Override AMapEntry nextResult (Tree tree) {
return tree;
}
@Override protected boolean isAfterIntendedEnd (Tree tree) {
return comparator.compare (tree.key, to) > 0;
}
};
}
else {
return new TreeIterator> (pathStack, first) {
@Override AMapEntry nextResult (Tree tree) {
return tree;
}
@Override protected boolean isAfterIntendedEnd (Tree tree) {
return comparator.compare (tree.key, to) >= 0;
}
};
}
}
@Override public Iterable> rangeII (final K fromKey, final K toKey) {
return new Iterable> () {
@Override public Iterator> iterator () {
return ARedBlackTreeMap.this.iterator (fromKey, toKey, true, true);
}
};
}
@Override public Iterable> rangeIE (final K fromKey, final K toKey) {
return new Iterable> () {
@Override public Iterator> iterator () {
return ARedBlackTreeMap.this.iterator (fromKey, toKey, true, false);
}
};
}
@Override public Iterable> rangeEI (final K fromKey, final K toKey) {
return new Iterable> () {
@Override public Iterator> iterator () {
return ARedBlackTreeMap.this.iterator (fromKey, toKey, false, true);
}
};
}
@Override public Iterable> rangeEE (final K fromKey, final K toKey) {
return new Iterable> () {
@Override public Iterator> iterator () {
return ARedBlackTreeMap.this.iterator (fromKey, toKey, false, false);
}
};
}
@Override public Iterable> fromI (final K fromKey) {
return new Iterable> () {
@Override public Iterator> iterator () {
return ARedBlackTreeMap.this.iterator (fromKey, null, true, false);
}
};
}
@Override public Iterable> fromE (final K fromKey) {
return new Iterable> () {
@Override public Iterator> iterator () {
return ARedBlackTreeMap.this.iterator (fromKey, null, false, false);
}
};
}
@Override public Iterable> toI (final K toKey) {
return new Iterable> () {
@Override public Iterator> iterator () {
return ARedBlackTreeMap.this.iterator (null, toKey, false, true);
}
};
}
@Override public Iterable> toE (final K toKey) {
return new Iterable> () {
@Override public Iterator> iterator () {
return ARedBlackTreeMap.this.iterator (null, toKey, false, false);
}
};
}
public void dump() {
dump (root, 0);
}
static void indent (int level) {
System.out.print (" ".substring (0, 4*level));
}
static void dump (Tree tree, int indent) {
if (tree == null) {
indent (indent);
System.out.println ("<>");
}
else {
dump (tree.left, indent+1);
indent (indent);
System.out.print (tree instanceof BlackTree ? "+ " : "* ");
System.out.println (tree.key);
dump (tree.right, indent+1);
}
}
static void validate(Tree tree) {
if (tree == null) {
return;
}
validate (tree.left);
validate (tree.right);
// rule 4: every redden node has two blacken children
if (isRedTree (tree)) {
if (tree.left != null && isRedTree (tree.left)) {
throw new IllegalStateException ("tree " + tree.key + " is redden and has a left child that is redden");
}
if (tree.right != null && isRedTree (tree.right)) {
throw new IllegalStateException ("tree " + tree.key + " is redden and has a right child that is redden");
}
}
checkBlackDepth (tree);
}
private static int checkBlackDepth (Tree tree) {
if (tree == null) {
return 1;
}
final int own = isBlackTree (tree) ? 1 : 0;
final int left = checkBlackDepth (tree.left);
final int right = checkBlackDepth (tree.right);
// rule 5: every path to 'leaf' nodes must have the same number of blacken nodes
if (left != right) {
throw new IllegalStateException ("left and right side have paths to leaf nodes with different numbers of blacken nodes: " + tree.key);
}
return own + left;
}
private abstract class TreeIterator implements Iterator {
@SuppressWarnings ("unchecked")
private final ArrayDeque> pathStack;
private Tree next;
abstract R nextResult (Tree tree); //TODO rename this
@SuppressWarnings ("unchecked")
private TreeIterator() {
// initialize 'next' with the leftmost element
if (root == null) {
pathStack = null;
next = null;
}
else {
pathStack = new ArrayDeque<> ();
next = root;
while (next.left != null) {
pathStack.push (next);
next = next.left;
}
}
}
private TreeIterator (ArrayDeque> pathStack, Tree first) {
this.pathStack = pathStack;
this.next = first;
}
@Override public boolean hasNext() {
return next != null;
}
@Override public R next() {
if (next == null) {
throw new NoSuchElementException ();
}
final Tree cur = next;
next = filteredFindNext (next.right);
return nextResult (cur);
}
@Override public void remove () {
throw new UnsupportedOperationException ();
}
private Tree filteredFindNext (Tree tree) {
final Tree result = findNext (tree);
if (result == null || isAfterIntendedEnd (result)) return null;
return result;
}
private Tree findNext (Tree tree) {
while (true) {
if (tree == null) {
return popPath ();
}
if (tree.left == null) {
return tree;
}
pathStack.push (tree);
tree = tree.left;
}
}
/**
* override if an iterator should finish before the entire map is traversed
*/
protected boolean isAfterIntendedEnd (Tree tree) {
return false;
}
private Tree popPath() {
if (pathStack.isEmpty ()) {
// convenience for handling the end of iteration
return null;
}
return pathStack.pop ();
}
}
static Tree lookup(Tree tree, K key, Comparator comparator) {
while (tree != null) {
final int cmp = comparator.compare (key, tree.key);
if (cmp == 0) return tree;
tree = (cmp < 0) ? tree.left : tree.right;
}
return null;
}
static boolean isRedTree (Tree tree) {
return tree != null && tree.isRed ();
}
static boolean isBlackTree (Tree tree) {
return tree != null && tree.isBlack ();
}
static Tree blacken (Tree tree) {
if (tree == null) {
return null;
}
return tree.blacken ();
}
static Tree balanceLeft (TreeFactory treeFactory, K key, V value, Tree l, Tree d) {
if (isRedTree (l) && isRedTree (l.left)) {
return new RedTree<> (l.key, l.value,
new BlackTree<> (l.left.key, l.left.value, l.left.left, l.left.right),
new BlackTree<> (key, value, l.right, d));
}
if (isRedTree (l) && isRedTree (l.right)) {
return new RedTree<> (l.right.key, l.right.value,
new BlackTree<> (l.key, l.value, l.left, l.right.left),
new BlackTree<> (key, value, l.right.right, d));
}
return treeFactory.create (l, d);
}
static Tree balanceRight (TreeFactory treeFactory, K key, V value, Tree a, Tree r) {
if (isRedTree (r) && isRedTree (r.left)) {
return new RedTree<> (r.left.key, r.left.value,
new BlackTree<> (key, value, a, r.left.left),
new BlackTree<> (r.key, r.value, r.left.right, r.right));
}
if (isRedTree (r) && isRedTree (r.right)) {
return new RedTree<> (r.key, r.value,
new BlackTree<> (key, value, a, r.left),
new BlackTree<> (r.right.key, r.right.value, r.right.left, r.right.right));
}
return treeFactory.create (a, r);
}
static Tree upd (Tree tree, K key, V value, Comparator comparator) {
if (tree == null) {
return new RedTree<> (key, value, null, null);
}
final int cmp = comparator.compare (key, tree.key);
if (cmp < 0) {
return balanceLeft (tree, tree.key, tree.value, upd (tree.left, key, value, comparator), tree.right);
}
if (cmp > 0) {
return balanceRight (tree, tree.key, tree.value, tree.left, upd (tree.right, key, value, comparator));
}
return tree.withNewValue (key, value);
}
static Tree del (Tree tree, K key, Comparator comparator) {
if (tree == null) {
return null;
}
final int cmp = comparator.compare(key, tree.key);
if (cmp < 0) {
// the node that must be deleted is to the left
return isBlackTree (tree.left) ?
balanceLeft (tree.key, tree.value, del (tree.left, key, comparator), tree.right) :
// tree.left is 'redden', so its children are guaranteed to be blacken.
new RedTree<> (tree.key, tree.value, del (tree.left, key, comparator), tree.right);
}
else if (cmp > 0) {
// the node that must be deleted is to the right
return isBlackTree (tree.right) ?
balanceRight (tree.key, tree.value, tree.left, del (tree.right, key, comparator)) :
new RedTree<> (tree.key, tree.value, tree.left, del (tree.right, key, comparator));
}
// delete this node and we are finished
return append (tree.left, tree.right);
}
static Tree balance (K key, V value, Tree tl, Tree tr) {
if (isRedTree (tl) && isRedTree (tr)) return new RedTree<> (key, value, tl.blacken (), tr.blacken ());
if (isRedTree (tl)) {
// left is redden, right is blacken
if (isRedTree (tl.left)) return new RedTree<> (tl.key, tl.value, tl.left.blacken (), new BlackTree<> (key, value, tl.right, tr));
if (isRedTree (tl.right)) {
return new RedTree<> (tl.right.key, tl.right.value,
new BlackTree<> (tl.key, tl.value, tl.left, tl.right.left),
new BlackTree<> (key, value, tl.right.right, tr));
}
return new BlackTree<> (key, value, tl, tr);
}
if (isRedTree (tr)) {
// left is blacken, right is redden
if (isRedTree (tr.right)) return new RedTree<> (tr.key, tr.value, new BlackTree<> (key, value, tl, tr.left), tr.right.blacken ());
if (isRedTree (tr.left)) return new RedTree<> (tr.left.key, tr.left.value, new BlackTree<> (key, value, tl, tr.left.left), new BlackTree<> (tr.key, tr.value, tr.left.right, tr.right));
return new BlackTree<> (key, value, tl, tr);
}
// tl and tr are both blacken
return new BlackTree<> (key, value, tl, tr);
}
private static Tree balanceLeft (K key, V value, Tree tl, Tree tr) { //TODO merge with other 'balanceLeft' method?
if (isRedTree (tl)) {
return new RedTree<> (key, value, tl.blacken (), tr);
}
if (isBlackTree (tr)) {
return balance (key, value, tl, tr.redden ());
}
if (isRedTree (tr) && isBlackTree (tr.left)) {
return new RedTree<> (tr.left.key, tr.left.value, new BlackTree<> (key, value, tl, tr.left.left), balance (tr.key, tr.value, tr.left.right, tr.right.blackToRed ()));
}
throw new IllegalStateException ("invariant violation");
}
static Tree balanceRight (K key, V value, Tree tl, Tree tr) {
if (isRedTree (tr)) {
return new RedTree<> (key, value, tl, tr.blacken ());
}
if (isBlackTree (tl)) {
return balance (key, value, tl.redden (), tr);
}
if (isRedTree (tl) && isBlackTree (tl.right)) {
return new RedTree<> (tl.right.key, tl.right.value, balance (tl.key, tl.value, tl.left.blackToRed (), tl.right.left), new BlackTree <> (key, value, tl.right.right, tr));
}
throw new IllegalStateException ("invariant violation");
}
/**
* This method combines two separate sub-trees into a single (balanced) tree. It assumes that both subtrees are
* balanced and that all elements in 'tl' are smaller than all elements in 'tr'. This situation occurs when a
* node is deleted and its child nodes must be combined into a resulting tree.
*/
private static Tree append (Tree tl, Tree tr) {
if (tl == null) return tr;
if (tr == null) return tl;
if (isRedTree (tl) && isRedTree (tr)) {
final Tree bc = append (tl.right, tr.left);
return isRedTree (bc) ?
new RedTree<> (bc.key, bc.value, new RedTree<> (tl.key, tl.value, tl.left, bc.left), new RedTree<> (tr.key, tr.value, bc.right, tr.right)) :
new RedTree<> (tl.key, tl.value, tl.left, new RedTree<> (tr.key, tr.value, bc, tr.right));
}
if (isBlackTree (tl) && isBlackTree (tr)) {
final Tree bc = append (tl.right, tr.left);
return isRedTree (bc) ?
new RedTree<> (bc.key, bc.value, new BlackTree<> (tl.key, tl.value, tl.left, bc.left), new BlackTree<> (tr.key, tr.value, bc.right, tr.right)) :
balanceLeft (tl.key, tl.value, tl.left, new BlackTree<> (tr.key, tr.value, bc, tr.right));
}
if (isRedTree (tr)) {
return new RedTree<> (tr.key, tr.value, append (tl, tr.left), tr.right);
}
if (isRedTree (tl)) {
return new RedTree<> (tl.key, tl.value, tl.left, append (tl.right, tr));
}
throw new IllegalStateException ("invariant violation: unmatched tree on append: " + tl + ", " + tr);
}
/**
* encapsulates tree creation for a given colour
*/
interface TreeFactory {
Tree create (Tree left, Tree right);
}
static abstract class Tree implements TreeFactory, AMapEntry {
final K key;
final V value;
final int count;
final Tree left;
final Tree right;
public Tree (K key, V value, Tree left, Tree right) {
this.key = key;
this.value = value;
this.left = left;
this.right = right;
this.count = 1 +
(left == null ? 0 : left.count) +
(right == null ? 0 : right.count);
}
@Override public K getKey () {
return key;
}
@Override public V getValue () {
return value;
}
abstract Tree withNewValue (K key, V value);
abstract Tree blackToRed();
abstract boolean isRed();
abstract boolean isBlack();
abstract Tree redden ();
abstract Tree blacken ();
}
static class BlackTree extends Tree {
public BlackTree (K key, V value, Tree left, Tree right) {
super(key, value, left, right);
}
@Override Tree withNewValue (K key, V value) {
return new BlackTree<> (key, value, left, right);
}
@Override public Tree create (Tree left, Tree right) {
return new BlackTree<> (key, value, left, right);
}
@Override Tree blackToRed () {
return redden ();
}
@Override boolean isRed () {
return false;
}
@Override boolean isBlack () {
return true;
}
@Override Tree redden () {
return new RedTree<> (key, value, left, right);
}
@Override Tree blacken () {
return this;
}
}
static class RedTree extends Tree {
public RedTree (K key, V value, Tree left, Tree right) {
super (key, value, left, right);
}
@Override Tree withNewValue (K key, V value) {
return new RedTree<> (key, value, left, right);
}
@Override public Tree create (Tree left, Tree right) {
return new RedTree<> (key, value, left, right);
}
@Override Tree blackToRed () {
throw new IllegalStateException ();
}
@Override boolean isRed () {
return true;
}
@Override boolean isBlack () {
return false;
}
@Override Tree redden () {
return this;
}
@Override Tree blacken () {
return new BlackTree<> (key, value, left, right);
}
}
}