The unexamined life is not worth living.
― Socrates
Contents
1. Introduction
A java Set is a Collection containing only unique elements. A TreeSet is an implementation of the Set which uses a Red-Black tree implementation. In this article, we demonstrate how to properly create and use a TreeSet.
2. Elements Must Be Comparable
The TreeSet sorts and stores the elements in the tree. It requires a Comparator function to sort the elements while storing. If one is not specified, it uses the elements natural ordering, which means the elements must implement Comparable. Note that in the following example, the TreeSet output is already sorted.
TreeSet<String> aset = new TreeSet<>();
aset.add("apple");
aset.add("orange");
aset.add("banana");
aset.add("grapes");
aset.add("melon");
System.out.println("t1: " + aset);
# prints
t1: [apple, banana, grapes, melon, orange]
3. Using a Comparator
When the elements you are adding does not implement the Comparable interface, then you need to use a Comparator function when creating the TreeSet. Here is a class which does not implement the Comparable interface.
public class User {
private String firstName;
private String lastName;
private int age;
public User(String firstName,String lastName,int age) {
this.firstName = firstName;
this.lastName = lastName;
this.age = age;
}
// standard getters and setters
public String toString()
{
StringBuilder sbuf = new StringBuilder();
sbuf.append('(').append(firstName)
.append(',').append(lastName)
.append(',').append(age)
.append(')');
return sbuf.toString();
}
}
And here, we use a comparison function comparing last names to store a bunch of users.
Comparator<User> c = (u1, u2) -> u1.getLastName().compareTo(u2.getLastName());
TreeSet<User> aset = new TreeSet<>(c);
aset.addAll(Arrays.asList(new User("Chris", "Pruitt", 34),
new User("Matt", "Bates", 15),
new User("John", "Wagner", 82),
new User("Vernon", "McGuire", 31),
new User("Mary", "Bates", 37),
new User("Mary", "Bates", 23)));
System.out.println("t2: " + aset);
# prints
t2: [(Matt,Bates,15), (Vernon,McGuire,31), (Chris,Pruitt,34), (John,Wagner,82)]
Note that, since we are comparing by last names, only the first entry with the last name Bates is added. The add() method (which addAll() uses) does not add the entry if one is already present.
Here is the result of using a comparer function comparing just the first names.
Comparator<User> c = (u1, u2) -> u1.getFirstName().compareTo(u2.getFirstName());
TreeSet<User> aset = new TreeSet<>(c);
aset.addAll(Arrays.asList(new User("Chris", "Pruitt", 34),
new User("Matt", "Bates", 15),
new User("John", "Wagner", 82),
new User("Vernon", "McGuire", 31),
new User("Mary", "Bates", 37),
new User("Mary", "Bates", 23)));
System.out.println("t3: " + aset);
# prints
t3: [(Chris,Pruitt,34), (John,Wagner,82), (Mary,Bates,37), (Matt,Bates,15), (Vernon,McGuire,31)]
Again, since we are attempting to add two Marys, the second one is skipped.
And here, we use a more complete comparison function which compares all fields, one after another.
Comparator<User> c = (u1, u2) -> {
int r = u1.getLastName().compareTo(u2.getLastName());
if ( r != 0 ) return r;
r = u1.getFirstName().compareTo(u2.getFirstName());
if ( r != 0 ) return r;
return Integer.compare(u1.getAge(), u2.getAge());
};
And the output:
System.out.println("t4: " + aset);
# prints
t4: [(Mary,Bates,23), (Mary,Bates,37), (Matt,Bates,15), (Vernon,McGuire,31), (Chris,Pruitt,34), (John,Wagner,82)]
4. Iterate In Ascending or Descending Order
The TreeSet provides an iterator() method which returns an Iterator for retrieving the items in the ascending order.
Iterator<User> it = aset.iterator();
System.out.println("t5: ");
while (it.hasNext()) {
User user = it.next();
System.out.println(" " + user);
}
# prints
t5:
(Mary,Bates,23)
(Mary,Bates,37)
(Matt,Bates,15)
(Vernon,McGuire,31)
(Chris,Pruitt,34)
(John,Wagner,82)
And of course there is a method descendingIterator() which is useful for retrieving items in the reverse order.
Iterator<User> it = aset.descendingIterator();
System.out.println("t6: ");
...
# prints
t6:
(John,Wagner,82)
(Chris,Pruitt,34)
(Vernon,McGuire,31)
(Matt,Bates,15)
(Mary,Bates,37)
(Mary,Bates,23)
5. Finding Elements
To check whether a particular element exists within a TreeSet is quite easy. Just use contains() to check for the presence.
# prints
System.out.println("t7: " + aset.contains(new User("Vernon", "McGuire", 31)));
t7: true
Of more interest is extracting elements which match a condition. Let us say, we want to find all users between the ages of 30 and 40. While the TreeSet supports headSet(), tailSet() and subSet() method for extracting a few elements, these cannot be used since the comparator has already been set and we cannot specify a range for the age.
The solution is to use java 8 streams. We can easily stream the contents of the Set and filter for the elements we want, and collect the results in a Set (or even a List).
Set<User> res = aset
.stream()
.filter(u -> u.getAge() >= 30 && u.getAge() < 40)
.collect(Collectors.toSet());
System.out.println("t8: " + res);
# prints
t8: [(Chris,Pruitt,34), (Mary,Bates,37), (Vernon,McGuire,31)]
6. Deleting Elements
Removing elements from the TreeSet can easily be accomplished using the removeIf() method which accepts a Predicate function. In our example, we remove users whose last name is Bates as follows:
System.out.println("t9: " + aset);
aset.removeIf(u -> u.getLastName().equals("Bates"));
System.out.println("t10: " + aset);
# prints
t9: [(Mary,Bates,23), (Mary,Bates,37), (Matt,Bates,15), (Vernon,McGuire,31), (Chris,Pruitt,34), (John,Wagner,82)]
t10: [(Vernon,McGuire,31), (Chris,Pruitt,34), (John,Wagner,82)]
Review
We have now learnt the basics of creating and using a TreeSet. Not using a Comparator while creating the TreeSet requires the elements to implement Comparable. We can extract the elements in ascending or descending order according to the comparator. Filter elements using a java 8 lambda function and remove them using removeIf().