TreeMap ​

We already know that HashMap is a mapping table that trades space for time. Its implementation principle dictates that the keys are unordered, meaning that the order in which the keys are traversed in a HashMap is unpredictable (though each key will be traversed exactly once).

There is another type of Map that sorts the keys internally, which is called SortedMap. Note that SortedMap is an interface, and its implementation class is TreeMap.

       ┌───┐
       │Map│
       └───┘
         ▲
    ┌────┴─────┐
    │          │
┌───────┐ ┌─────────┐
│HashMap│ │SortedMap│
└───────┘ └─────────┘
               ▲
               │
          ┌─────────┐
          │ TreeMap │
          └─────────┘

SortedMap guarantees that the traversal is sorted in the order of the keys. For example, if the keys "apple," "pear," and "orange" are inserted, the traversal order will definitely be "apple," "orange," "pear," since String is sorted alphabetically by default:

import java.util.*;

public class Main {
    public static void main(String[] args) {
        Map<String, Integer> map = new TreeMap<>();
        map.put("orange", 1);
        map.put("apple", 2);
        map.put("pear", 3);
        for (String key : map.keySet()) {
            System.out.println(key);
        }
        // apple, orange, pear
    }
}

When using TreeMap, the keys must implement the Comparable interface. Classes like String and Integer already implement the Comparable interface, so they can be used directly as keys. There are no requirements for objects used as values.

If the class used as a key does not implement the Comparable interface, a custom sorting algorithm must be specified when creating the TreeMap:

import java.util.*;

public class Main {
    public static void main(String[] args) {
        Map<Person, Integer> map = new TreeMap<>(new Comparator<Person>() {
            public int compare(Person p1, Person p2) {
                return p1.name.compareTo(p2.name);
            }
        });
        map.put(new Person("Tom"), 1);
        map.put(new Person("Bob"), 2);
        map.put(new Person("Lily"), 3);
        for (Person key : map.keySet()) {
            System.out.println(key);
        }
        // {Person: Bob}, {Person: Lily}, {Person: Tom}
        System.out.println(map.get(new Person("Bob"))); // 2
    }
}

class Person {
    public String name;
    Person(String name) {
        this.name = name;
    }
    public String toString() {
        return "{Person: " + name + "}";
    }
}

Note that the Comparator interface requires the implementation of a comparison method, which compares two elements a and b. If a < b, it should return a negative number (usually -1); if a == b, it should return 0; if a > b, it should return a positive number (usually 1). The TreeMap sorts the keys based on the comparison results.

From the output of the code above, we can see that the printed keys are indeed sorted according to the order defined by the Comparator. To find a value based on a key, we can pass in new Person("Bob") as the key, and it will return the corresponding integer value 2.

Additionally, note that the Person class does not override equals() and hashCode(), as TreeMap does not use equals() and hashCode().

Now, let's look at a slightly more complex example: this time we define a Student class and sort it based on the score, with higher scores coming first:

import java.util.*;

public class Main {
    public static void main(String[] args) {
        Map<Student, Integer> map = new TreeMap<>(new Comparator<Student>() {
            public int compare(Student p1, Student p2) {
                return p1.score > p2.score ? -1 : 1;
            }
        });
        map.put(new Student("Tom", 77), 1);
        map.put(new Student("Bob", 66), 2);
        map.put(new Student("Lily", 99), 3);
        for (Student key : map.keySet()) {
            System.out.println(key);
        }
        System.out.println(map.get(new Student("Bob", 66))); // null?
    }
}

class Student {
    public String name;
    public int score;
    Student(String name, int score) {
        this.name = name;
        this.score = score;
    }
    public String toString() {
        return String.format("{%s: score=%d}", name, score);
    }
}

In the for loop, we do get the correct order. However, wait! When we look for the value using the same key, new Student("Bob", 66), the result is null!

What could be the problem? Is there an issue with TreeMap? When encountering issues with TreeMap, we should first review the basic principles of Java programming: when a problem arises, do not doubt the Java standard library; instead, look for the cause in your own code.

In this example, the problem with TreeMap actually stems from the Comparator:

public int compare(Student p1, Student p2) {
    return p1.score > p2.score ? -1 : 1;
}

When p1.score and p2.score are not equal, the return value is correct; however, when they are equal, it does not return 0! This is the reason why TreeMap is not working correctly: TreeMap relies on the compareTo() method of the keys or the Comparator.compare() method to determine if two keys are equal. It must return 0 when the two keys are equal. Therefore, modify the code as follows:

public int compare(Student p1, Student p2) {
    if (p1.score == p2.score) {
        return 0;
    }
    return p1.score > p2.score ? -1 : 1;
}

Alternatively, you can also use Integer.compare(int, int) to get the correct comparison results.

Summary ​

SortedMap strictly traverses the keys in order during iteration, and its most commonly used implementation class is TreeMap.

The keys of a SortedMap must implement the Comparable interface or provide a Comparator.

It is crucial to implement the comparison logic according to the specifications of compare(); otherwise, TreeMap will not function correctly.