TreeMap and TreeSet of a custom Java type


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:

  • 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
        }
    }
  • Another solution is to provide the TreeSet with a Comparator
  • Version ≥ 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:

  • 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.
  • It's important to implement the comparison properly and consistently.
    • 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.

Basic Programs