How to Add Custom Class Objects to the TreeSet in Java?

TreeSet is an implementation of the SortedSet interface in java that uses a red-black tree for storage. By default, It maintains an ascending order. It contains unique elements only. It doesn’t allow null elements. Access and retrieval times are quite fast. To add the user-defined object into TreeSet, we need to implement a Comparable interface. An element that we want to add in TreeSet must be a comparable type. If we don’t implement the Comparable interface then it will throw ClassCastException. 

Example 1

Java

// Importing util package
import java.util.*;

// Custom class Car implements Comparable interface
class Car implements Comparable<Car> {

    // attributes
    int Modelno;
    String name, city;
    int stock;

    // Car constructor
    public Car(int Modelno, String name, String city,
               int stock)
    {
        this.Modelno = Modelno;
        this.name = name;
        this.city = city;
        this.stock = stock;
    }

    // Override the compareTo() method
    public int compareTo(Car c)
    {
        if (stock > c.stock) {
            return 1;
        }
        else if (stock < c.stock) {
            return -1;
        }
        else {
            return 0;
        }
    }
}

public class GFG {

    // Main driver method
    public static void main(String[] args)
    {
        // Define an objects of TreeSet class
        TreeSet<Car> set = new TreeSet<Car>();

        // Creating Car objects
        Car c1 = new Car(132, "BMW", "Rajkot", 35);
        Car c2 = new Car(269, "Audi", "Surat", 20);
        Car c3 = new Car(560, "Kia", "Vadodara", 15);
        Car c4 = new Car(109, "Creta", "Ahmedabad", 26);

        // Adding Car objects to TreeSet
        set.add(c1);
        set.add(c2);
        set.add(c3);
        set.add(c4);

        // Traversing TreeSet
        for (Car c : set) {
            System.out.println(c.stock + " " + c.name + " "
                               + c.city + " " + c.Modelno);
        }
    }
}

Output

15 Kia Vadodara 560
20 Audi Surat 269
26 Creta Ahmedabad 109
35 BMW Rajkot 132

We need to override the compareTo() method so that it will sort our set in a particular order. 

Example 2: Let’s take the same example but now we override the compareTo() method with respect to the name of the car.

Java

// Importing util package
import java.util.*;

// Custom class Car implements Comparable interface
class Car implements Comparable<Car> {

    // attributes
    int Modelno;
    String name, city;
    int stock;

    // Car constructor
    public Car(int Modelno, String name, String city,
               int stock)
    {
        this.Modelno = Modelno;
        this.name = name;
        this.city = city;
        this.stock = stock;
    }

    // Override the compareTo() method
    public int compareTo(Car c)
    {
        return name.compareTo(c.name);
    }
}

public class GFG {
    // Main driver method
    public static void main(String[] args)
    {

        // Define an objects of TreeSet class
        TreeSet<Car> set = new TreeSet<Car>();

        // Creating Car objects
        Car c1 = new Car(132, "BMW", "Rajkot", 35);
        Car c2 = new Car(269, "Audi", "Surat", 20);
        Car c3 = new Car(560, "Kia", "Vadodara", 15);
        Car c4 = new Car(109, "Creta", "Ahmedabad", 26);

        // Adding Car objects to TreeSet
        set.add(c1);
        set.add(c2);
        set.add(c3);
        set.add(c4);

        // Traversing TreeSet
        for (Car c : set) {
            System.out.println(c.name + " " + c.stock + " "
                               + c.city + " " + c.Modelno);
        }
    }
}

Output

Audi 20 Surat 269
BMW 35 Rajkot 132
Creta 26 Ahmedabad 109
Kia 15 Vadodara 560