Left Leaning Red/Black Tree implemented in Java
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(); | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment