How to return the k-th element in TreeSet in Java?

Question

Maybe I am not using the right data structure. I need to use a set, but also want to efficiently return the k-th smallest element. Can TreeSet in Java do this? There seems no built-in method of TreeSet to do this.

Answer 1

I don't believe that TreeSet has a method that directly does this. There are binary search trees that do support O(log n) random access (they are sometimes called order statistic trees), and there are Java implementations of this data structure available. These structures are typically implemented as binary search trees that store information in each node counting how many elements are to the left or right of the node, so a search down the tree can be made to find the appropriate element by descending into the appropriate subtree at each step. The classic "Introduction to Algorithms, Third Edition" book by Cormen, Rivest, Leisserson, and Stein explores this data structure in their chapter "Augmenting Data Structures" if you are curious how to implement one yourself.

Alternatively, you may be able (in some cases) to use the TreeSet's tailSet method and a modified binary search to try to find the kth element.  Specifically, look at the first and last elements of the TreeSet, then (if possible given the contents) pick some element that is halfway between the two and pass it as an argument to tailSet to get a view of the elements of the set after the midpoint. Using the number of elements in the tailSet, you could then decide whether you've found the element, or whether to explore the left or right halves of the tree. This is a slightly modified interpolation search over the tree, and could potentially be fast. However, I don't know the internal complexity of the tailSet methods, so this could be actually be worse than the order statistic tree. It also might fail if you can't compute the "midpoint" of two elements, for example, if you are storing Strings in your TreeSet.

Answer 2

You just need to iterate to element k. One way to do that would be to use one of Guava's Iterables.get methods:

T element = Iterables.get(set, k);

There's no built in method to do this because a Set is not a List and index-based operations like that are generally the reserved for Lists. A TreeSet is more appropriate for things like finding the closest contained element that is >= some value.

One thing you could do if the fastest possible access to the kth smallest element were really important would be to use an ArrayList rather than a TreeSet and handle inserts by binary searching for the insertion point and either inserting the element at that index or replacing the existing element at that index, depending on the result of the search. Then you could get the kth smallest element in O(1) by just calling get(k).

You could even create an implementation of SortedSet that handles all that and adds the get(index) method if you really wanted.