Search code examples
javaalgorithmdata-structurestreesetsortedset

How to find the index of an element in a TreeSet?


I'm using a TreeSet<Integer> and I'd quite simply like to find the index of a number in the set. Is there a nice way to do this that actually makes use of the O(log(n)) complexity of binary trees?

(If not, what should I do, and does anyone know why not? I'm curious why such a class would be included in Java without something like a search function.)


Solution

  • As @Yrlec points out set.headSet(element).size will returns 0 though there is no this element in the set. So we'd better check:

     return set.contains(element)? set.headSet(element).size(): -1;
    

    Here is a test case to show the problem:

    public static void main(String args[]){
        TreeSet<Integer> set = new TreeSet<>();
        set.add(4);
        set.add(2);
        set.add(3);
        set.add(1);
    
        System.out.println(set.headSet(1).size());//0
        System.out.println(set.headSet(2).size());//1
        System.out.println(set.headSet(3).size());//2
        System.out.println(set.headSet(4).size());//3
        System.out.println(set.headSet(-1).size());//0!!Caution,returns 0 though it does not exist!
    
    }