Java Language TreeMap/TreeSet of a custom Java type


Example

Since TreeMaps and TreeSets maintain keys/elements according to their natural ordering. Therefor TreeMap keys and TreeSet elements have to comparable to one another.

Say we have a custom Person class:

public class Person  {

    private int id;
    private String firstName, lastName;
    private Date birthday;

    //... Constuctors, getters, setters and various methods
}

If we store it as-is in a TreeSet (or a Key in a TreeMap):

TreeSet<Person2> set = ...      
set.add(new Person(1,"first","last",Date.from(Instant.now())));

Then we'd run into an Exception such as this one:

Exception in thread "main" java.lang.ClassCastException: Person cannot be cast to java.lang.Comparable
    at java.util.TreeMap.compare(TreeMap.java:1294)
    at java.util.TreeMap.put(TreeMap.java:538)
    at java.util.TreeSet.add(TreeSet.java:255)

To fix that, let's assume that we want to order Person instances based on the order of their ids (private int id). We could do it in one of two ways:

  1. One solution is to modify Person so it would implement the Comparable interface:

    public class Person implements Comparable<Person> {
        private int id;
        private String firstName, lastName;
        private Date birthday;
    
        //... Constuctors, getters, setters and various methods
    
        @Override
        public int compareTo(Person o) {
            return Integer.compare(this.id, o.id); //Compare by id
        }
    }
    
  2. Another solution is to provide the TreeSet with a Comparator:

Java SE 8
TreeSet<Person> treeSet = new TreeSet<>((personA, personB) -> Integer.compare(personA.getId(), personB.getId()));
TreeSet<Person> treeSet = new TreeSet<>(new Comparator<Person>(){
    @Override
    public int compare(Person personA, Person personB) {
        return Integer.compare(personA.getId(), personB.getId());
    }
});

However, there are two caveats to both approaches:

  1. It's very important not to modify any fields used for ordering once an instance has been inserted into a TreeSet/TreeMap. In the above example, if we change the id of a person that's already inserted into the collection, we might run into unexpected behavior.

  2. It's important to implement the comparison properly and consistently. As per the Javadoc:

The implementor must ensure sgn(x.compareTo(y)) == -sgn(y.compareTo(x)) for all x and y. (This implies that x.compareTo(y) must throw an exception iff y.compareTo(x) throws an exception.)

The implementor must also ensure that the relation is transitive: (x.compareTo(y)>0 && y.compareTo(z)>0) implies x.compareTo(z)>0.

Finally, the implementor must ensure that x.compareTo(y)==0 implies that sgn(x.compareTo(z)) == sgn(y.compareTo(z)), for all z.