Why does not a remove call to this Treeset effectively removes the element?

Question

I was trying earlier to code a Dijkstra algorithm and faced out a problem with my implementation. I finally found a solultion but am still puzzled with the behaviour of the code.

What I am trying to do here is to iteratively mark the elements of a tree set. Everytime I mark an element, I remove it from the set and reinsert it, so the set places it at the end of the set (as the compareTo method tells it to do). Doing so, I can take the first element of the set inside a while loop and stop when this first element is marked .

The following example is a minimal example:

public static void main(final String[] args) {
    Set<Element> elements = new TreeSet<>();
    elements.add(new Element(1));
    elements.add(new Element(2));
    elements.add(new Element(3));
    elements.add(new Element(5));
    elements.add(new Element(4));
    System.err.println(elements.size());
    while(!elements.iterator().next().mark){
        Element element = elements.iterator().next();
        element.mark();
        elements.remove(element);
        elements.add(element);
    }
    System.err.println(elements.size());
}

private static class Element implements Comparable<Element>{

    private boolean mark = false;
    private final int id;

    public Element(int id){
        this.id = id;
    }

    public void mark(){
        this.mark = true;
    }

    @Override
    public boolean equals(Object o){
        Element other = (Element) o;
        if(this.mark == other.mark){
            return this.id == other.id;
        } else {
            return false;
        }
    }

    @Override
    public int compareTo(Element other){
        if(this.mark == other.mark){
            return this.id - other.id;
        } else if(this.mark && !other.mark){
            return 1;
        } else {
            return -1;
        }
    }
}

The problem here is that the line :

elements.remove(element);

Does not remove the element. However, running it in a debugger, I can see that the element has the exact same reference as the element in the set and is equal to it with respect to the equals method. So, when running the line :

elements.add(element);

The same element is added two times ; because I suppose that the insertion in a tree set uses a binary search and does not test the equality between the first element of the set and the newly inserted one.

I have understood that swapping the two lines :

element.mark();
elements.remove(element);

Actually solves the problem.

My question though is : why does a modified element of a tree set is not effectively removed when modified through an iterator ?

Answer 1

TreeSet uses compareTo in order to position the elements in a tree structure (and also to determine equality). If you put an element in the TreeSet and then modify a property that is used in your compareTo method (ie the mark in your case), remove can't locate that element, since it looks for it in a different location than the one in which the original element was inserted.

If you remove the element before changing it, you solve the problem, since the original element is found were it was inserted.