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.

Solution 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.

Solution 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.

Solution 3

Use TreeSet.iterator() to get an iterator in ascending order and call next() K times:

// Example for Integers
Iterator<Integer> it = treeSet.iterator();
int i = 0;
Integer current = null;
while(it.hasNext() && i < k) {
   current = it.next();
   i++;
}

Solution 4

https://github.com/geniot/indexed-tree-map

I had the same problem. So I took the source code of java.util.TreeMap and wrote IndexedTreeMap. It implements my own IndexedNavigableMap:

public interface IndexedNavigableMap<K, V> extends NavigableMap<K, V> {
   K exactKey(int index);
   Entry<K, V> exactEntry(int index);
   int keyIndex(K k);
}

The implementation is based on updating node weights in the red-black tree when it is changed. Weight is the number of child nodes beneath a given node, plus one - self. For example when a tree is rotated to the left:

    private void rotateLeft(Entry<K, V> p) {
    if (p != null) {
        Entry<K, V> r = p.right;

        int delta = getWeight(r.left) - getWeight(p.right);
        p.right = r.left;
        p.updateWeight(delta);

        if (r.left != null) {
            r.left.parent = p;
        }

        r.parent = p.parent;


        if (p.parent == null) {
            root = r;
        } else if (p.parent.left == p) {
            delta = getWeight(r) - getWeight(p.parent.left);
            p.parent.left = r;
            p.parent.updateWeight(delta);
        } else {
            delta = getWeight(r) - getWeight(p.parent.right);
            p.parent.right = r;
            p.parent.updateWeight(delta);
        }

        delta = getWeight(p) - getWeight(r.left);
        r.left = p;
        r.updateWeight(delta);

        p.parent = r;
    }
  }

updateWeight simply updates weights up to the root:

   void updateWeight(int delta) {
        weight += delta;
        Entry<K, V> p = parent;
        while (p != null) {
            p.weight += delta;
            p = p.parent;
        }
    }

And when we need to find the element by index here is the implementation that uses weights:

public K exactKey(int index) {
    if (index < 0 || index > size() - 1) {
        throw new ArrayIndexOutOfBoundsException();
    }
    return getExactKey(root, index);
}

private K getExactKey(Entry<K, V> e, int index) {
    if (e.left == null && index == 0) {
        return e.key;
    }
    if (e.left == null && e.right == null) {
        return e.key;
    }
    if (e.left != null && e.left.weight > index) {
        return getExactKey(e.left, index);
    }
    if (e.left != null && e.left.weight == index) {
        return e.key;
    }
    return getExactKey(e.right, index - (e.left == null ? 0 : e.left.weight) - 1);
}

Also comes in very handy finding the index of a key:

    public int keyIndex(K key) {
    if (key == null) {
        throw new NullPointerException();
    }
    Entry<K, V> e = getEntry(key);
    if (e == null) {
        throw new NullPointerException();
    }
    if (e == root) {
        return getWeight(e) - getWeight(e.right) - 1;//index to return
    }
    int index = 0;
    int cmp;
    if (e.left != null) {
        index += getWeight(e.left);
    }
    Entry<K, V> p = e.parent;
    // split comparator and comparable paths
    Comparator<? super K> cpr = comparator;
    if (cpr != null) {
        while (p != null) {
            cmp = cpr.compare(key, p.key);
            if (cmp > 0) {
                index += getWeight(p.left) + 1;
            }
            p = p.parent;
        }
    } else {
        Comparable<? super K> k = (Comparable<? super K>) key;
        while (p != null) {
            if (k.compareTo(p.key) > 0) {
                index += getWeight(p.left) + 1;
            }
            p = p.parent;
        }
    }
    return index;
}

You can find the result of this work at https://github.com/geniot/indexed-tree-map.

Solution 5

TreeSet<Integer> a=new TreeSet<>();
a.add(1);
a.add(2);
a.add(-1);

System.out.println(a.toArray()[0]);

it can be helpfull

View more solutions

Comments

[James Montagne]

Though the original question is about HashMap, this question has some discussion of TreeMap as well. stackoverflow.com/questions/2656995/nth-item-of-hashmap

[templatetypedef]

This works, but it's extremely slow (O(k) time). Perhaps there is a faster approach?

[Tudor]

@templatetypedef: Ah yes, you are right. I missed the 'efficient' requirement.

[templatetypedef]

But this runs in O(k) time, which isn't particularly efficient.

[Admin]

"Java implementations of this data structure" link does not seem to be correct. :)

[Admin]

I think I provided enough information. I said I want to use a set with operations finding kth smallest element. So whatever standard set operations + finding kth smallest element.

[Admin]

btw, I need to use set, because I need not recount duplicates.

[Louis Wasserman]

It's as efficient as you're going to get with a Java TreeSet, I'm afraid, at least in the general case.

[ColinD]

For the record, getting the size of a tailSet or headSet of a TreeSet requires iterating through the elements and counting them. So your suggestion about that probably isn't going to help.

[templatetypedef]

@ColinD- That sounds like an implementation detail to me; I could easily see an implementation that supports size in O(1).

[noooooooob]

In TreeSet API, both k and k + 1 are referring to the actually element in the set, rather than index. So this does not work here.

[aclowkay]

It doesn't scale very well. But In my case it worked great, since my treeset is usually small.