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);
    }
}

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]

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);
    }
}

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]

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));
    }
}

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

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));
    }
}

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

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

• Java Fundamentals

• Array Interview Essentials

• Java Memory Essentials

• Java Keywords Essentials

• Java OOPs Essentials

• Java Strings Essentials

Happy Coding!