How to return the k-th element in TreeSet in Java?
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
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.