Order statistic tree in Java

Hello everyone ,
I need to use a balanced tree implementation which supports retrieve rank of a given entry in O ( log n ) time . Currently the way I see to do this would be to use TreeSet which is a red-black tree implementation . Then call headMap() method of the given entry which runs in O(1) and then call size() method on the headMap() received which takes O(size of head map) time ( hence linear ) . I need to be able to do this in log(N) . Obviously it is possible to implement a tree on my own which takes care of rank issue also . Such balanced trees which support rank queries in O ( log N ) are called order statistic trees . My question is doesn’t there exist something in the JAVA API itself which can help me do this .

2 Likes

@everyone : Figured out the answer . There is no default implementation in java library for this . I had to modify Java’s TreeSet or TreeMap class to contain an extra field and extra function for implementing an order-statistic tree . Got the AC in contest timings itself . Cool :slight_smile:

Hey , some people were asking me a link to the implementation .
Here is the solution which uses it .
http://www.codechef.com/viewsolution/1819919 ( Solution Page )
http://www.codechef.com/FEB13/problems/MBOARD/ ( Related Problem Page , problem figured in Feb 2013 contest ) .
Here you can find the class MyTreeMap as an inner class which is an implementation of order-statistic tree in Java .

1 Like

I added a couple methods to vinnetpaliwal’s code to look up the element at rank i and output all the keys to an array.

public K getKeyForRank(int i) {
	Entry<K, V> e = root;
	for (;;) {
		int leftSize = e.left == null ? 0 : e.left.size;
		if (leftSize <= i) {
			i -= leftSize;
			if (i > 0) {
				i--;
				e = e.right;
			} else
				break;
		} else
			e = e.left;
	}
	return e.key;
}

public K[] keysToArray(K[] a) {
	if (a.length < size)
		a = (K[]) java.lang.reflect.Array.newInstance(a.getClass()
				.getComponentType(), size);
	addToArray(root, a, 0);
	return a;
}

private void addToArray(Entry<K, V> e, K[] a, int i) {
	if (e == null)
		return;
	addToArray(e.left, a, i);
	int leftSize = e.left == null ? 0 : e.left.size;
	a[i + leftSize] = e.key;
	addToArray(e.right, a, i + leftSize + 1);
}