RedBlackTree.java

package treehack;

import java.util.Comparator;
import java.util.Iterator;
import java.util.TreeSet;

public class RedBlackTree<T> implements Iterable<T> {
	// java 8 guarantees that TreeSet uses red black trees
	private TreeSet<T> treeSet;

	public RedBlackTree() {
		treeSet = new TreeSet<>();
	}

	public RedBlackTree(Comparator<T> comparator) {
		treeSet = new TreeSet<>(comparator);
	}

	public T get(int rank) {
		return treeSet.stream().skip(rank).findFirst().orElse(null);
	}

	public boolean insert(T element) {
		return treeSet.add(element);
	}

	public boolean remove(T element) {
		return treeSet.remove(element);
	}

	public int size() {
		return treeSet.size();
	}

	@Override
	public Iterator<T> iterator() {
		return treeSet.iterator();
	}
}