package org.aspyct.going;
import java.util.Iterator;
import java.util.Stack;
/**
* This is an implementation of the Left Leaning Red/Black Tree as
* described in the course from the University of Princeton.
*
* This code contains a lot of assertions (assert ;).
* They are included only to make things as clear as possible.
* These can be safely removed.
*
* @author Antoine d'Otreppe <a.dotreppe@aspyct.org> @aspyct
*
*/
public class RBTree<K extends Comparable<K>, V> implements Iterable<RBTree<K, V>.Entry<K, V>>{
private static final boolean RED = true;
private static final boolean BLACK = false;
public abstract class Entry<K, V> {
public abstract K key();
public abstract V value();
}
private class Node<K, V> extends Entry<K, V> {
K key;
V value;
Node<K, V> left;
Node<K, V> right;
boolean color;
Node(K key, V value) {
this.key = key;
this.value = value;
this.color = RED;
}
public K key() {
return key;
}
public V value() {
return value;
}
}
private Node<K, V> root;
public void put(K key, V value) {
root = put(root, key, value);
root.color = BLACK;
}
private Node<K, V> put(Node<K, V> node, K key, V value) {
if (node == null) {
return new Node<K, V>(key, value);
}
// Inv: the current node is not a 4-node
if (isFourNode(node)) {
node = splitFourNode(node);
}
assert !isFourNode(node);
// Regular recursive BST insert
int cmp = key.compareTo(node.key);
if (cmp == 0) {
node.value = value;
}
else if (cmp < 0) {
node.left = put(node.left, key, value);
}
else /* if (cmp > 0) */ {
node.right = put(node.right, key, value);
}
// No right-leaning red branch please
if (isRed(node.right)) {
node = leanLeft(node);
}
assert !isRed(node.right);
if (node.left != null) {
assert node.left.key.compareTo(node.key) < 0;
}
if (node.right != null) {
assert node.right.key.compareTo(node.key) > 0;
}
assert node != null;
return node;
}
private boolean isFourNode(Node<K, V> node) {
return isRed(node.left) && isRed(node.left.left);
}
private boolean isRed(Node<K, V> node) {
if (node == null) {
return false;
}
else {
return node.color == RED;
}
}
private Node<K, V> splitFourNode(Node<K, V> node) {
/**
* Scenario: node is black
* node.left is red
* node.left.left is red
*/
assert !isRed(node);
assert isRed(node.left);
assert isRed(node.left.left);
node = rotateRight(node);
node.left.color = BLACK;
return node;
}
private Node<K, V> leanLeft(Node<K, V> node) {
/**
* Scenario: node.right is red
*/
assert isRed(node.right);
node = rotateLeft(node);
node.color = node.left.color;
node.left.color = RED;
return node;
}
private Node<K, V> rotateLeft(Node<K, V> node) {
Node<K, V> top = node.right;
node.right = top.left;
top.left = node;
return top;
}
private Node<K, V> rotateRight(Node<K, V> node) {
Node<K, V> top = node.left;
node.left = top.right;
top.right = node;
return top;
}
public Iterator<Entry<K, V>> iterator() {
return new InorderEntryIterator(root);
}
private class InorderEntryIterator implements Iterator<Entry<K, V>> {
Stack<Node<K, V>> stack;
InorderEntryIterator(Node<K, V> root) {
stack = new Stack<Node<K, V>>();
explore(root);
}
private void explore(Node<K, V> node) {
while (node != null) {
stack.push(node);
node = node.left;
}
}
public boolean hasNext() {
return !stack.empty();
}
public Entry<K, V> next() {
Node<K, V> node = stack.pop();
if (node.right != null) {
explore(node.right);
}
return node;
}
public void remove() {
throw new UnsupportedOperationException("Cannot remove nodes from here");
}
}
// java -ea org.aspyct.going.RBTree
public static void main(String[] args) {
RBTree<Character, Integer> tree = new RBTree<Character, Integer>();
/*
* c
* / \
*/
tree.put('c', 0);
assert tree.root != null;
assert tree.root.key == 'c';
assert tree.root.value == 0;
assert tree.root.color == BLACK;
/* Take a bigger key and check that the tree is rotated
*
* i
* / \
* c
* / \
*/
tree.put('i', 0);
assert tree.root.key == 'i';
assert tree.root.left.key == 'c';
assert tree.root.color == BLACK;
assert tree.root.left.color == RED;
assert tree.root.right == null;
assert tree.root.left.right == null;
assert tree.root.left.left == null;
/* Again, a bigger key, and check that the tree is still left
*
* v
* / \
* i
* / \
* c
* / \
*/
tree.put('v', 0);
assert tree.root.key == 'v';
assert tree.root.left.key == 'i';
assert tree.root.left.left.key == 'c';
assert tree.root.color == BLACK;
assert tree.root.left.color == RED;
assert tree.root.left.left.color == RED;
assert tree.root.right == null;
assert tree.root.left.right == null;
assert tree.root.left.left.right == null;
assert tree.root.left.left.left == null;
/* Last check, the four node must be split
* Try inserting something between i and v
*
*
* i
* / \
* c v
* / \ / \
* m
* / \
*/
tree.put('m', 0);
assert tree.root.key == 'i';
assert tree.root.right.key == 'v';
assert tree.root.left.key == 'c';
assert tree.root.right.left.key == 'm';
assert tree.root.color == BLACK;
assert tree.root.left.color == BLACK;
assert tree.root.right.color == BLACK;
assert tree.root.right.left.color == RED;
assert tree.root.left.left == null;
assert tree.root.left.right == null;
assert tree.root.right.left.left == null;
assert tree.root.right.left.right == null;
assert tree.root.right.right == null;
System.out.println("All good :)");
System.out.println("Inorder iterate");
for (RBTree<Character, Integer>.Entry<Character, Integer> entry: tree) {
System.out.print(entry.key().toString() + " ");
}
System.out.println();
}
}
