How TreeMap can save your day?

In Java collections, the TreeMap is a sorted and navigable map that organizes elements in a self-balancing binary tree. It can help you solve problems like “Which element in this collection is the closest to this one?”

Introduction

You want to plan reservations for a room. Here are some reservations:

• From 4-Jan to 7-Jan, occupant: Mr. A
• From 10-Jan to 21-Jan, occupant: Mrs. B
• From 5-Feb to 18-Feb, occupant: Mr. A
• From 20-Feb to 3-Mar, occupant: Mr. C

Suppose that there is no possibility for a reservation to override another one. How do you determine in a programmatic way who is the occupant the 10-Feb? And the 19-Feb?

Algorithm

There are many possibilities to answer to those questions. Basically, one of them is to store the reservations in a set. When searching for an occupant at a given date, you walk through the set, testing each element. You stop when you have a reservation that contains the given date or when there is no more reservation to test. This is a linear search algorithm. In  the worst case, the time complexity is O(n), ie. when you have to test all elements and there is no matching element or the matching element is the last one.

To improve this algorithm, you can also use a binary search algorithm. It consists at each step in dividing the set of elements in two parts and continuing searching in one of those parts. This algorithm needs a set of sorted elements. Its time complexity is O(log(n)) in the worst case. This is better because if you have 7 elements, you will only need 3 tests against 7 for linear search, in the worst case. And if you have 1000 elements, you will only need 10 tests against 1000 for the linear version, in the worst case. However, with the binary search algorithm, you may lost time while sorting elements. Java proposes in classes java.util.Arrays and java.util.Collections the method sort() the uses the merge sort algorithm that guaranties a time complexity of O(n.log(n)). Thus, sorting + binary search is worst than a linear search, that doesn’t need a sorted set of elements. Note that Java SE proposes implementations of the binary search algorithm in the methods java.util.Arrays.binarySearch() and java.util.Collections.binarySearch().

Another approach to the binary search is to generate the corresponding binary tree. Indeed, the binary search algorithm is like a walk through a height-balanced binary tree. At each step of the algorithm, we select the left part (left branch) or right part (right branch) of the subset of elements (the subtree) according to a central element (a node). But the big difference is that if you apply an algorithm like the red-black tree for your binary tree, two aspects are guaranteed:

1. Each time you add an element, it is sorted with the rest of the tree.
2. The tree is height-balanced. It means that its maximal height is O(log(n)).

For such a tree, the operations insert and search cost of time is O(log(n)) each.

The class java.util.TreeMap represents such a height-balancing binary tree and uses the red-black tree algorithm. But what is even more interesting with this class, is that it provides the methods ceilingEntry() and floorEntry(). They return the map entry which key is the closest respectively after and before the given key.

Representation of a reservation

Below is the source code of a simplified class to represent such reservations. This implementation uses guava.

public class Reservation {
    public Date from;
    public Date to;
    public String occupant;

    public Reservation(Date from, Date to, String occupant) {
        this.from = from;
        this.to = to;
        this.occupant = occupant;
    }

    public boolean contains(Date date) {
        return !(from.after(date) || to.before(date));
    }

    @Override
    public String toString() {
        return Objects.toStringHelper(this)
            .add("from", from)
            .add("to", to)
            .add("occupant", occupant)
            .toString();
    }

    @Override
    public int hashCode() {
        return Objects.hashCode(from, to, occupant);
    }

    @Override
    public boolean equals(Object obj) {
        if (this == obj) {
            return true;
        }
        if (obj == null || !(obj instanceof Reservation)) {
            return false;
        }
        Reservation reservation = (Reservation) obj;
        return Objects.equals(from, reservation.from)
            && Objects.equals(to, reservation.to)
            && Objects.equals(occupant, reservation.occupant);
    }

}

We declare below all four reservations that you can find in the introduction at the top of the post.

Date from, to;
Calendar calendar = Calendar.getInstance();

calendar.set(2011, 0, 4);
from = calendar.getTime();
calendar.set(2011, 0, 7);
to = calendar.getTime();
Reservation reserv1MrA = new Reservation(from, to, "Mr. A");

calendar.set(2011, 0, 10);
from = calendar.getTime();
calendar.set(2011, 0, 21);
to = calendar.getTime();
Reservation reserv1MrsB = new Reservation(from, to, "Mrs. B");

calendar.set(2011, 1, 5);
from = calendar.getTime();
calendar.set(2011, 1, 18);
to = calendar.getTime();
Reservation reserv2MrA = new Reservation(from, to, "Mr. A");

calendar.set(2011, 1, 20);
from = calendar.getTime();
calendar.set(2011, 2, 3);
to = calendar.getTime();
Reservation reserv1MrC = new Reservation(from, to, "Mr. C");

How to find a reservation from a given date

I group a set of reservations in a class that I’ve called Planning. Once instantiated, you can add reservations (method add()) and you can get a reservation at a given date (method getReservationAt()).

There are two steps to get a reservation close to a date.

1. I first use the method floorEntry(). It gets you the map entry which key is the greatest one less than the given date.
2. Second, I check if the reservation (that is the value of the entry) contains the given date.

public class Planning {
    public final TreeMap<Date, Reservation> reservations;

    public Planning() {
        reservations = new TreeMap<Date, Reservation>();
    }

    public void add(Reservation reservation) {
        reservations.put(reservation.from, reservation);
    }

    public Reservation getReservationAt(Date date) {
        Entry<Date, Reservation> entry = reservations.floorEntry(date);
        if (entry == null) {
            return null;
        }
        Reservation reservation = entry.getValue();
        if (!reservation.contains(date)) {
            return null;
        }
        return reservation;
    }
}

Test

We first store the reservations declared above in a planning.

Planning planning = new Planning();
planning.add(reserv1MrA);
planning.add(reserv1MrsB);
planning.add(reserv2MrA);
planning.add(reserv1MrC);

Now, who has made a reservation the 10-Feb and the 19-Feb? (Here, I use FEST-assert as assertion API.)

Calendar calendar = Calendar.getInstance();

calendar.set(2011, 1, 10);
Date dateOfMrA = calendar.getTime();
assertThat(planning.getReservationAt(dateOfMrA).occupant).isEqualTo("Mr. A");

calendar.set(2011, 1, 19);
Date dateNoReservation = calendar.getTime();
assertThat(planning.getReservationAt(dateNoReservation)).isNull();