TreeMap in Java

The TreeMap class implements the Map interface and is part of the Java Collection Interface. Unlike other Map implementations, a TreeMap stores the key-value pairs in sorted order.

In this tutorial, we will learn more about the sorting functionality of TreeMap.

Natural Ordering in TreeMap

A TreeMap, by default, stores the entries sorted by the keys. Keys are sorted by using their natural ordering. For example, the natural ordering for integers is the ascending order. Similarly, the natural ordering for strings is the alphabetical order.

Let's create a TreeMap that stores integer-string pairs. When we print the map, the output is displayed in sorted order.

import java.util.Map.Entry;
import java.util.TreeMap;

public class TreeMapDemo
{
	public static void main(String[] args)
	{
		TreeMap<Integer, String> map = new TreeMap<>();
		map.put(4, "Harry");
		map.put(3, "Simon");
		map.put(2, "Jessica");
		map.put(5, "Victor");
		map.put(1, "Justin");
		
		//Printing the TreeMap
		for(Entry<Integer, String> e : map.entrySet())
			System.out.println(e.getKey() + "-->" + e.getValue());
	}
}

1-->Justin
2-->Jessica
3-->Simon
4-->Harry
5-->Victor

Let's use the strings as keys and view the output.

import java.util.Map.Entry;
import java.util.TreeMap;

public class TreeMapDemo
{
	public static void main(String[] args)
	{
		TreeMap<String, Integer> map = new TreeMap<>();
		map.put("Harry", 4);
		map.put("Jessica", 2);
		map.put("Victor", 5);
		map.put("Simon", 3);
		map.put("Justin", 1);
		
		//Printing the TreeMap
		for(Entry<String, Integer> e : map.entrySet())
			System.out.println(e.getKey() + "-->" + e.getValue());
	}
}

Harry-->4
Jessica-->2
Justin-->1
Simon-->3
Victor-->5

The Comparable interface is used to set the natural order for user-defined classes. Let's create a Student class that implements the Comparable interface. We need to sort the Students by their names(alphabetical order). TreeMap will automatically use the natural order defined in the compareTo() method to sort and store the data.

import java.util.Map.Entry;
import java.util.TreeMap;

class Student implements Comparable
{
	String name;
	Double gpa;
	
	Student(String s, Double d)
	{
		this.name = s;
		this.gpa = d;
	}

	//Defining the natural ordering for Students
	@Override
	public int compareTo(Object o)
	{
		return this.name.compareTo(( (Student) o).name);
	}
}
public class TreeMapDemo
{
	public static void main(String[] args)
	{
		TreeMap<Student, String> map = new TreeMap<>() ;
		map.put((new Student("Harry", 4.0)), "Harry");
		map.put((new Student("Jessica", 2.8)), "Jessica");
		map.put((new Student("Victor", 3.9)), "Victor");
		map.put((new Student("Simon", 4.4)), "Simon");
		map.put((new Student("Justin",3.01)), "Justin");
		
		//Printing the TreeMap
		for(Entry<Student, String> e : map.entrySet())
			System.out.println(e.getKey().gpa + "-->" + e.getValue());
	}
}

4.0-->Harry
2.8-->Jessica
3.01-->Justin
4.4-->Simon
3.9-->Victor

Defining a Different Order for Sorting

The natural ordering may not be ideal in some cases. The TreeMap constructor can take a Comparator as an argument. We can use this Comparator to define a different order of sorting keys.

For example, we can use the reverse-order Comparator if we want to store the keys in descending order.

import java.util.Comparator;
import java.util.Map.Entry;
import java.util.TreeMap;

public class TreeMapDemo
{
	public static void main(String[] args)
	{
        //Passing reverse order Comparator
		TreeMap<String, Integer> map = new TreeMap<>(Comparator.reverseOrder());
		map.put("Harry", 4);
		map.put("Jessica", 2);
		map.put("Victor", 5);
		map.put("Simon", 3);
		map.put("Justin", 1);
		
		//Printing the TreeMap
		for(Entry<String, Integer> e : map.entrySet())
			System.out.println(e.getKey() + "-->" + e.getValue());
	}
}

Victor-->5
Simon-->3
Justin-->1
Jessica-->2
Harry-->4

Let's try to create a custom Comparator to work with the Student class defined earlier. Instead of sorting the Students by their names, we will sort them based on their GPA. A student with a higher GPA should come first in the TreeMap.

import java.util.Comparator;
import java.util.Map.Entry;
import java.util.TreeMap;

class Student
{
	String name;
	Double gpa;
	
	Student(String s, Double d)
	{
		this.name = s;
		this.gpa = d;
	}
}
public class TreeMapDemo
{
	public static void main(String[] args)
	{
		//Comparator to sort by descending order of GPA
        Comparator<Student> c = (Student s1, Student s2) -> (-s1.gpa.compareTo(s2.gpa));		
		TreeMap<Student, String> map = new TreeMap<>(c) ;
		map.put((new Student("Harry", 4.0)), "Harry");
		map.put((new Student("Jessica", 2.8)), "Jessica");
		map.put((new Student("Victor", 3.9)), "Victor");
		map.put((new Student("Simon", 4.4)), "Simon");
		map.put((new Student("Justin",3.01)), "Justin");
		
		//Printing the TreeMap
		for(Entry<Student, String> e : map.entrySet())
			System.out.println(e.getKey().gpa + "-->" + e.getValue());
	}
}

4.4-->Simon
4.0-->Harry
3.9-->Victor
3.01-->Justin
2.8-->Jessica

How Is TreeMap Implemented in Java?

• TreeMap implements the NavigableMap interface. It uses an advanced data structure called the Red-Black Trees to store and sort the key-value pairs.
• Red-Black trees are very similar to normal Binary Search Trees. Just like a Binary Search Tree, a Red-Black tree will contain a root, a left subtree, and a right subtree. All nodes in the left subtree should be smaller than the node, and all nodes of the right tree should be greater than the node.
• This property should be true for all other nodes. The natural ordering of the objects determines the smaller and greater objects.
• Unlike a normal BST, a Red-Black tree is self-balancing. This makes sure that no one half of the tree becomes a lot deeper than the other.
• In short, this property gives a time complexity of O(logN) for the basic operations on the tree. BST, on the other hand, can give a worst-case time complexity of O(N).

Common TreeMap Methods

The TreeMap class provides additional methods to take advantage of the sorted order of keys. Let's take a look at a few of the important methods of the TreeMap class.

• The firstKey() method is used to fetch the first or the smallest key from the TreeMap.
• Similarly, the lastKey() method is used to fetch the largest key from the map.
• The headMap() method takes a key as input. It returns a map that contains the key-value pairs smaller than the key passed.
• Similarly, the tailMap() method takes a key and returns a map containing all key-value pairs greater than and equal to the passed key.

import java.util.Map.Entry;
import java.util.Set;
import java.util.TreeMap;

public class TreeMapDemo
{
	public static void main(String[] args)
	{
		TreeMap<String, Integer> map = new TreeMap<>();
		map.put("Harry", 4);
		map.put("Jessica", 2);
		map.put("Victor", 5);
		map.put("Simon", 3);
		map.put("Justin", 1);
		System.out.println("The TreeMap is: ");
		for(Entry<String, Integer> e : map.entrySet())
			System.out.println(e.getKey() + "-->" + e.getValue());
		
		String smallestKey = map.firstKey();
		String largestKey = map.lastKey();
		Set<String> keysSmallerThanVictor = map.headMap("Victor").keySet();//Excludes Victor
		Set<String> keysGreaterThanJustin = map.tailMap("Justin").keySet();//Includes Justin
		
		System.out.println("\nSmallest Key: " + smallestKey);
		System.out.println("Largest Key: " + largestKey);
		System.out.println("Keys Smaller than Victor: " + keysSmallerThanVictor);
		System.out.println("Keys Greater than Justin: " + keysGreaterThanJustin);
	}
}

The TreeMap is:
Harry-->4
Jessica-->2
Justin-->1
Simon-->3
Victor-->5

Smallest Key: Harry
Largest Key: Victor
Keys Smaller than Victor: [Harry, Jessica, Justin, Simon]
Keys Greater than Justin: [Justin, Simon, Victor]

When to use TreeMap?

• A HashMap maintains no order of insertion and is used more as a general-purpose map. Storing and fetching data is very efficient and fast.
• A LinkedHashMap is also pretty fast and also maintains the order of insertion.
• A TreeMap is preferred over the HashMap and LinkedHashMap when we want to store key-value pairs in sorted order. We get complete control over the order in which data is stored in TreeMap.
• However, the sorting takes a toll on the overall performance of the TreeMap.
• It all boils down to the task which we are trying to accomplish. If performance is crucial, then prefer HashMap or LinkedHashMap. If the data must be stored in sorted order, then TreeMap should be preferred.

Summary

TreeMap is an important Java collection. TreeMap uses the Red-Black tree. This will make sure that insertion, deletion, and searching of elements take logarithmic time. But its performance is not as good as the HashMap or the LinkedHashMap. The TreeMap class provides a lot of flexibility in sorting the key-value pairs. We can either use the natural ordering of keys or define our own order using Comparators. In this tutorial, we discussed the basics of the TreeMap class and how to use it effectively.