Store unique objects in sorted order using TreeSet in Java

In Java, the TreeSet is a special implementation of Set interface to store unique objects in sorted order. By default it stores the objects in ascending order by performing simple compare. There are couple of ways to use the TreeSet as per our requirement. Here we would discuss about few most commonly used ways. First the simple implementation, then using Comparator and Comparable interface.



1. Storing simple objects

import java.util.*;

class TreeTestSimple {
    public static void main(String [] argv) {
        Set<String> set = new TreeSet<String>();
        set.add("Gabriel");
        set.add("Vikash");
        set.add("Josh");
        set.add("Macy");
        set.add("Liam");
        
        for (String data : set) {
            System.out.println(data);
        }
    }
}

Output:

Gabriel
Josh
Liam
Macy
Vikash


2. Storing complex objects

In reality, our objects are not that simple. You’d like to store objects of complex class like:

class Student {
    int rollNo;
    String name;
    String section;
    int totalMarks;
}

And say you want this to be sorted in order of rollNo.


2.1. Implementing Comparable

Easiest way is implement the Comparable interface into Student class:

import java.util.*;

import java.util.*;

class Student implements Comparable<Student> {
    Integer rollNo;
    String name;
    String section;
    int totalMarks;

    public Student(int rollNo, String name) {
        this.rollNo = rollNo;
        this.name = name;
    }

    @Override
    public String toString() {
        return String.format("%d: %s", rollNo, name);
    }

    @Override
    public int compareTo(Student o) {
        return this.rollNo.compareTo(o.rollNo);
    }
}

class TreeTestComplex1 {

    public static void main(String [] argv) {

        Set<Student> students = new TreeSet<Student>();

        students.add(new Student(54, "Prakash"));
        students.add(new Student(21, "Johny"));
        students.add(new Student(47, "Sunil"));
        students.add(new Student(10, "Blanton"));

        for (Student student : students) {
            System.out.println(student);
        }
    }
}

Output

10: Blanton
21: Johny
47: Sunil
54: Prakash

If you want to store in reverse order you can just rewrite the compareTo(...) method like

    @Override
    public int compareTo(Student o) {
        return o.rollNo.compareTo(this.rollNo);
    }


2.2. Using custom Comparator while constructing TreeSet

There can be situation where you want to store the Student objects in more than one Set in order of different fields. Better way to do this using Comparator instead of implementing Comparable. You need to pass new Comparator implementation while constructing the TreeSet.

import java.util.*;

class Student {
    Integer rollNo;
    String name;
    String section;
    int totalMarks;

    public Student(int rollNo, String name) {
        this.rollNo = rollNo;
        this.name = name;
    }

    @Override
    public String toString() {
        return String.format("%d: %s", rollNo, name);
    }
}

class TreeTestComplex2 {

    public static void main(String [] argv) {

        Set<Student> studentsByRollNo = new TreeSet<Student>(
                new Comparator<Student>() {
                    @Override
                    public int compare(Student o1, Student o2) {
                        return o1.rollNo.compareTo(o2.rollNo);
                    }
                }
        );

        Set<Student> studentsByName = new TreeSet<Student>(
                new Comparator<Student>() {
                    @Override
                    public int compare(Student o1, Student o2) {
                        return o1.name.compareTo(o2.name);
                    }
                }
        );

        Student student1 = new Student(54, "Prakash"),
                student2 = new Student(23, "Adrea"),
                student3 = new Student(47, "Sunil"),
                student4 = new Student(10, "Blanton");


        studentsByRollNo.add(student1);
        studentsByRollNo.add(student2);
        studentsByRollNo.add(student3);
        studentsByRollNo.add(student4);


        studentsByName.add(student1);
        studentsByName.add(student2);
        studentsByName.add(student3);
        studentsByName.add(student4);


        System.out.println("Students by Roll No:");
        for (Student student : studentsByRollNo) {
            System.out.println(student);
        }

        System.out.println("\nStudents by Name:");
        for (Student student : studentsByName) {
            System.out.println(student);
        }
    }
}

Output

Students by Roll No:
10: Blanton
23: Adrea
47: Sunil
54: Prakash

Students by Name:
23: Adrea
10: Blanton
54: Prakash
47: Sunil


For more information, please refer: - https://docs.oracle.com/javase/8/docs/api/java/util/TreeSet.html - https://docs.oracle.com/javase/8/docs/api/java/lang/Comparable.html - https://docs.oracle.com/javase/8/docs/api/java/util/Comparator.html


Hansaraj avatar
About Hansaraj
Hansaraj is a Software Engineer experienced in Java, Groovy, JavaScript, SQL, C#, C++
comments powered by Disqus