/ Java  

Java TreeSet

The TreeSet in Java is a subclass of Set. The TreeSet collection is used to sort the object elements, and it can also guarantee the uniqueness of the elements.

TreeSet is an implementation class of SortedSet interface. SortedSet can ensure that the collection elements are in sorted state.



The underlying structure of TreeSet is a black red tree. Every new element (except the first one) will be compared with the last inserted element and arrange them according to the structure of the binary tree.

There are two comparative methods in TreeSet. They are Comparable interface and Comparator interface.

Using Comparable

Using Comparable interface requires Object to implement Comparable interface. The Comparable interface requres the implementation of compareTo(Object object) method. TreeSet will call the compareTo(Object object) method of elements to compare the relationship between the elements, and then arranges the collection elements in ascending order.

If the Object doesn’t implement Comparable interface, a ClassCastException will be thrown.

Some common classes in Java have implemented the Comparable interface: BigDecimal, BigInteger, Character, Boolean, String, Date, Time, etc.

Example:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
public static void main(String[] args) {
    TreeSet<Person> ts = new TreeSet<>();
    
    ts.add(new Person("Person A", 23));
    ts.add(new Person("Person B", 44));
    ts.add(new Person("Person C", 13));
    ts.add(new Person("Person D", 43));
    ts.add(new Person("Person E", 33));
        
    System.out.println(ts);
}

public class Person implements Comparable<Person> {
    private String name;
    private int age;

    // Compare by age
    public int compareTo(Person o) {
        return this.age - o.age;      
    }
}

// Output
// [
//  Person{name='Person C', age=13},
//  Person{name='Person A', age=23},
//  Person{name='Person E', age=33},
//  Person{name='Person D', age=43},
//  Person{name='Person B', age=44}
// ]

When the compareTo method returns 0, the two elements are considered to be equal. Therefore, when we need to put an object into the TreeSet and rewrite the equals() method of this class, we should ensure that when the equals() method returns true, the compareTo method returns 0.

After an element have been added to the TreeSet, it’s not recommended to modify the value of the elements. It is easy to cause some errors. Example:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
public static void main(String[] args) {
    TreeSet<Person> ts = new TreeSet<>();

    ts.add(new Person("Person A", 23));
    ts.add(new Person("Person B", 44));
    ts.add(new Person("Person C", 13));
    ts.add(new Person("Person D", 43));
    ts.add(new Person("Person E", 33));

    ts.forEach(obj -> System.out.println("Before modification: " + obj.getName() + " " + obj.getAge()));

    Person person = ts.first();
    person.setAge(80);

    person = ts.last();
    person.setAge(5);

    ts.forEach(obj -> System.out.println("After modification: " + obj.getName() + " " + obj.getAge()));
}

TreeSet will not be re-sorted after we modify the values. Also, if we make the elements equal after the modification, the TreeSet will not be re-sorted as well. At this time, if we want to delete the modified element, it will fail, and the element equal to the modified element will also be deleted. However, we can delete elements that have not changed normally.

Using Comparator

We can also define a comparator to implement the Comparator interface, rewrite the compare method, and pass the comparator into the TreeSet constructor.

Example:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
public class ComparatorDemo {
    public static void main(String[] args) {
        TreeSet<PersonB> ts = new TreeSet<>(new Comparator<PersonB>() {
            @Override
            public int compare(PersonB o1, PersonB o2) {
                return Integer.compare(o1.age, o2.age);
            }
        });

        ts.add(new PersonB( 23));
        ts.add(new PersonB(44));
        ts.add(new PersonB(13));

        ts.forEach(System.out::println);
    }

    static class PersonB {
        int age;

        public PersonB(int age) {
            this.age = age;
        }

        @Override
        public String toString() {
            return "Age: " + age;
        }
    }
}

Conclusion

  1. Feature
    • TreeSet is used for sorting, you can specify an order. After the object is stored, it will be arranged in the specified order.
  2. Usage
    • Comparable
      • Call the compareTo() method of the object and compare it with the objects in the collection
      • Store the elements according to the result returned by the compareTo() method
    • Comparator
      • A Comparator can be specified when creating a TreeSet
      • If a subclass object of Comparator is passed in, the TreeSet will be sorted according to the order specified by the comparator
      • The add() method will automatically call the compare() method of the Comparator interface to sort
      • The called object is the first parameter of the compare() method, and the object in the collection is the second parameter of the compare() method
    • Difference
      • If the TreeSet constructor gets nothing, it will by default use Comparable (ClassCastException is thrown if the object doesn’t implement Comparable)
      • If TreeSet constructor gets a Comparator, it will be given priority to Comparator