DEV Community

Cover image for Sorting Smarts in Java: TreeSet and TreeMap
Arshi Saxena
Arshi Saxena

Posted on

3 2 1 1 2

Sorting Smarts in Java: TreeSet and TreeMap

In our previous post, we discussed Comparable and Comparator, which are the building blocks for sorting elements in Java.

Now that we understand the fundamentals, let’s explore how collections like TreeSet and TreeMap leverage these interfaces to maintain elements in a sorted order.


TreeSet: A Set That Sorts

TreeSet is part of Java’s Set collection. Unlike HashSet, which does not guarantee any order, TreeSet stores its elements in a sorted order. It relies on the Comparable or Comparator implementation to determine how elements should be ordered.

Key Characteristics of TreeSet:

  • Sorted Order: Automatically maintains elements in sorted order.
  • Unique Elements: Does not allow duplicate elements.
  • Implementation: Based on a Red-Black Tree.

Example 1: Natural Ordering with TreeSet

// Implementing Comparable Interface
class Product implements Comparable<Product> {
    private int id;
    private double price;
    private String category;

    public Product(int id, double price, String category) {
        this.id = id;
        this.price = price;
        this.category = category;
    }

    // Overriding compareTo()
    @Override
    public int compareTo(Product o) {
        // Natural ordering by id
        return Integer.compare(this.id, o.id);
    }

    @Override
    public String toString() {
        return "Product [id=" + id + ", price=" + price + ", category=" +
        category + "]";
    }

    // Getters
    public int getId() {
        return id;
    }

    public double getPrice() {
        return price;
    }

    public String getCategory() {
        return category;
    }
}

public class TreeSetExample {
    public static void main(String[] args) {
        Set<Product> treeSet = new TreeSet<>();
        treeSet.add(new Product(3, 200.0, "Non-Essentials"));
        treeSet.add(new Product(1, 100.0, "Essentials"));
        treeSet.add(new Product(2, 150.0, "Essentials"));

        System.out.println("TreeSet Sorted with Natural Ordering:");
        treeSet.forEach(System.out::println);
    }
}
Enter fullscreen mode Exit fullscreen mode

Output:

TreeSet Sorted with Natural Ordering (by ID):
Product [id=1, price=100.0, category=Essentials]
Product [id=2, price=150.0, category=Essentials]
Product [id=3, price=200.0, category=Non-Essentials]
Enter fullscreen mode Exit fullscreen mode

Explanation:

  • The TreeSet sorts the products based on their id in ascending order because we defined the natural order using the Comparable interface.

  • Duplicate products with the same id would be ignored due to the nature of Set.


Example 2: Custom Ordering with TreeSet

We can also specify a custom sorting order using a Comparator. For instance, let’s sort by category and price:

public class CustomTreeSetExample {
    public static void main(String[] args) {
        // Define the Comparator for custom ordering
        Comparator<Product> productComparator = (p1, p2) -> {
            // Comparing on the basis of category
            int categoryComparison = p1.getCategory()
                                        .compareTo(p2.getCategory());

            // Comparing on the basis of price when categories are same
            return categoryComparison != 0 
                   ? categoryComparison 
                   : Double.compare(p1.getPrice(), p2.getPrice());
        };

        // Create TreeSet with the custom comparator
        Set<Product> treeSet = new TreeSet<>(productComparator);

        treeSet.add(new Product(3, 200.0, "Non-Essentials"));
        treeSet.add(new Product(1, 100.0, "Essentials"));
        treeSet.add(new Product(2, 250.0, "Essentials"));

        System.out.println(
            "TreeSet Sorted with Custom Ordering (by Category and Price):"
        );
        treeSet.forEach(System.out::println);
    }
}
Enter fullscreen mode Exit fullscreen mode

Output:

TreeSet Sorted with Custom Ordering (by Category and Price):
Product [id=1, price=100.0, category=Essentials]
Product [id=2, price=250.0, category=Essentials]
Product [id=3, price=200.0, category=Non-Essentials]
Enter fullscreen mode Exit fullscreen mode

Explanation:

The TreeSet now sorts first by the category and then by price within the same category using the custom comparator.


TreeMap: A Map That Sorts

TreeMap is similar to TreeSet, but it is a Map that stores key-value pairs in a sorted order. Like TreeSet, it uses the Comparable or Comparator interface to sort the keys.

Key Characteristics of TreeMap:

  • Sorted Keys: Stores keys in a sorted order.
  • Unique Keys: Does not allow duplicate keys.
  • Implementation: Based on a Red-Black Tree.

Example 1: Natural Ordering with TreeMap

// Implementing Comparable Interface
class Product implements Comparable<Product> {
    private int id;
    private double price;
    private String category;

    public Product(int id, double price, String category) {
        this.id = id;
        this.price = price;
        this.category = category;
    }

    // Overriding compareTo() method
    @Override
    public int compareTo(Product o) {
        // Natural ordering by id
        return Integer.compare(this.id, o.id);
    }

    @Override
    public String toString() {
        return "Product [id=" + id + ", price=" + price + ", category=" +
        category + "]";
    }

    // Getters
    public int getId() {
        return id;
    }

    public double getPrice() {
        return price;
    }

    public String getCategory() {
        return category;
    }
}

public class TreeMapExample {
    public static void main(String[] args) {
        // Initializing TreeMap with Product as key and Quantity as value
        Map<Product, Integer> productMap = new TreeMap<>();

        productMap.put(new Product(3, 200.0, "Non-Essentials"), 5);
        productMap.put(new Product(1, 100.0, "Essentials"), 10);
        productMap.put(new Product(2, 150.0, "Essentials"), 8);

        // Displaying the TreeMap
        System.out.println("TreeMap with Natural Ordering (Product by
        ID):");
        productMap.forEach((product, qty) -> System.out.println(product +
        " | Quantity: " + qty));
    }
}
Enter fullscreen mode Exit fullscreen mode

Output

TreeMap with Natural Ordering (Product by ID):
Product [id=1, price=100.0, category=Essentials] | Quantity: 10
Product [id=2, price=150.0, category=Essentials] | Quantity: 8
Product [id=3, price=200.0, category=Non-Essentials] | Quantity: 5
Enter fullscreen mode Exit fullscreen mode

Code Explanation

  1. Product Class:
    The Product class implements the Comparable interface to define the natural ordering based on the id.

  2. TreeMap:
    The TreeMap stores Product objects as keys, and the sorting order is determined by the compareTo() method of the Product class. As you can see, the TreeMap arranges the products based on their id in ascending order.

  3. Sorting in Action:
    When we insert the products into the TreeMap, it automatically sorts them by their id. The quantities are stored as values, and each Product is mapped to a specific quantity.


TreeMap: Custom Ordering with Comparator

Now, let’s use a custom comparator to sort the products in a TreeMap based on their category and price. This allows us to define a sorting logic that doesn’t rely on the natural ordering by id.

Example: TreeMap with Custom Ordering (By Category and Price)

public class TreeMapCustomOrdering {
    public static void main(String[] args) {
        // Custom comparator for sorting by category and price
        Comparator<Product> customComparator = (p1, p2) -> {
            // Comparing on the basis of category
            int categoryComparison = p1.getCategory()
                           .compareTo(p2.getCategory());

            // Comparing on the basis of price when categories are same
            return categoryComparison != 0 
                   ? categoryComparison 
                   : Double.compare(p1.getPrice(), p2.getPrice());
        };

        // Initializing TreeMap with custom comparator
        Map<Product, Integer> customSortedMap = new TreeMap<>(customComparator);

        customSortedMap.put(new Product(1, 100.0, "Essentials"), 10);
        customSortedMap.put(new Product(2, 500.0, "Essentials"), 15);
        customSortedMap.put(new Product(3, 200.0, "Non-Essentials"), 5);
        customSortedMap.put(new Product(4, 400.0, "Non-Essentials"), 7);
        customSortedMap.put(new Product(5, 300.0, "Essentials"), 8);

        // Displaying the TreeMap
        System.out.println("TreeMap with Custom Ordering (By Category and
        Price):");
        customSortedMap.forEach((product, qty) ->
        System.out.println(product + " | Quantity: " + qty));
    }
}
Enter fullscreen mode Exit fullscreen mode

Output

TreeMap with Custom Ordering (By Category and Price):
Product [id=1, price=100.0, category=Essentials] | Quantity: 10
Product [id=5, price=300.0, category=Essentials] | Quantity: 8
Product [id=2, price=500.0, category=Essentials] | Quantity: 15
Product [id=3, price=200.0, category=Non-Essentials] | Quantity: 5
Product [id=4, price=400.0, category=Non-Essentials] | Quantity: 7
Enter fullscreen mode Exit fullscreen mode

Code Explanation

  1. Custom Comparator:
    The custom comparator first compares the category fields. If the categories are the same, it then compares the price fields. This allows sorting by multiple criteria: first by category, and if the categories are the same, then by price.

  2. TreeMap with Comparator:
    The TreeMap is initialized with the custom comparator, which sorts the products accordingly. The products are first sorted by their category (alphabetically), and if two products have the same category, they are sorted by price in ascending order.

  3. Dynamic Sorting:
    This approach provides more flexibility than natural ordering, as we can define any comparison logic we need, including sorting by multiple attributes.


Conclusion

In this post, we explored how collections like TreeSet and TreeMap maintain sorted order using natural ordering and custom ordering via Comparable and Comparator. These tree-based collections are powerful tools when you need to keep elements in order without manually sorting them every time.


Key Takeaways:

  • TreeSet: A sorted set based on natural or custom order, with unique elements.
  • TreeMap: A sorted map that stores key-value pairs, with keys sorted in natural or custom order.

By combining the flexibility of Comparable and Comparator with these collections, you can efficiently manage sorted data structures in your Java programs.


Related Posts

Happy Coding!

Reinvent your career. Join DEV.

It takes one minute and is worth it for your career.

Get started

Top comments (0)

Great read:

Is it Time to go Back to the Monolith?

History repeats itself. Everything old is new again and I’ve been around long enough to see ideas discarded, rediscovered and return triumphantly to overtake the fad. In recent years SQL has made a tremendous comeback from the dead. We love relational databases all over again. I think the Monolith will have its space odyssey moment again. Microservices and serverless are trends pushed by the cloud vendors, designed to sell us more cloud computing resources.

Microservices make very little sense financially for most use cases. Yes, they can ramp down. But when they scale up, they pay the costs in dividends. The increased observability costs alone line the pockets of the “big cloud” vendors.

👋 Kindness is contagious

Immerse yourself in a wealth of knowledge with this piece, supported by the inclusive DEV Community—every developer, no matter where they are in their journey, is invited to contribute to our collective wisdom.

A simple “thank you” goes a long way—express your gratitude below in the comments!

Gathering insights enriches our journey on DEV and fortifies our community ties. Did you find this article valuable? Taking a moment to thank the author can have a significant impact.

Okay