How to write a custom Comparator for TreeMap in Java?

Question

I want to store key-value pairs in TreeMap and sort the entries based on the value of Key as per following logic:

Sort by the length of the key. If the length of two keys is same then sort them alphabetically. Example, for the following key-value pairs.

IBARAKI MitoCity
TOCHIGI UtunomiyaCity
GUNMA MaehashiCity
SAITAMA SaitamaCity
CHIBA ChibaCity
TOKYO Sinjyuku
KANAGAWA YokohamaCity

The expected output is like this.

CHIBA : ChibaCity
GUNMA : MaehashiCity
TOKYO : Sinjyuku
IBARAKI : MitoCity
SAITAMA : SaitamaCity
TOCHIGI : UtunomiyaCity
KANAGAWA : YokohamaCity

Solution 1

You need to write your own comparator for this and use it in TreeMap, e.g.:

public class StringComparator implements Comparator<String> {

    @Override
    public int compare(String s1, String s2) {
        return s1.length() == s2.length() ? s1.compareTo(s2) : s1.length() - s2.length();
    }

    public static void main(String[] args) throws JsonParseException, JsonMappingException, IOException {
        Map<String, String> map = new TreeMap<>(new StringComparator());
        map.put("IBARAKI", "MitoCity");
        map.put("TOCHIGI", "UtunomiyaCity");
        map.put("GUNMA", "MaehashiCity");
        map.put("SAITAMA", "SaitamaCity");
        map.put("CHIBA", "ChibaCity");
        map.put("TOKYO", "Sinjyuku");
        map.put("KANAGAWA", "YokohamaCity");

        System.out.println(map);
    }

}

This does not handle null values but you can add the handling if you are expecting null values in your use case.

Solution 2

You can pass the Comparator as a parameter to Map's constructor. According to documentation it is used for Keys only:

/**
 * Constructs a new, empty tree map, ordered according to the given
 * comparator.  All keys inserted into the map must be <em>mutually
 * comparable</em> by the given comparator: {@code comparator.compare(k1,
 * k2)} must not throw a {@code ClassCastException} for any keys
 * {@code k1} and {@code k2} in the map.  If the user attempts to put
 * a key into the map that violates this constraint, the {@code put(Object
 * key, Object value)} call will throw a
 * {@code ClassCastException}.
 *
 * @param comparator the comparator that will be used to order this map.
 *        If {@code null}, the {@linkplain Comparable natural
 *        ordering} of the keys will be used.
 */
public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}

In this way you can pass comparator by length of your key like this:

new TreeMap<>(Comparator.comparingInt(String::length).thenComparing(Comparator.naturalOrder()))

Solution 3

You can do this as follows.

  public static void main(String[] args) {

      Map<String, String> map = new TreeMap<>(new CustomSortComparator());

      map.put("IBARAKI", "MitoCity");
      map.put("TOCHIGI", "UtunomiyaCity");
      map.put("GUNMA", "MaehashiCity");
      map.put("SAITAMA", "SaitamaCity");
      map.put("CHIBA", "ChibaCity");
      map.put("TOKYO", "Sinjyuku");
      map.put("KANAGAWA", "YokohamaCity");

      System.out.println(map);

  }

The CustomSortComparator has been defined as follows.

public class CustomSortComparator implements Comparator<String> {

  @Override
  public int compare(String o1, String o2) {
    if (o1.length() > o2.length()) {
      return 1;
    }
    if (o1.length() < o2.length()) {
      return -1;
    }
    return returnCompareBytes(o1, o2);
  }

  private int returnCompareBytes(String key1, String key2) {
    for (int i = 0; i < key1.length() - 1; i++) {
      if (key1.charAt(i) > key2.charAt(i)) {
        return 1;
      }
      if (key1.charAt(i) < key2.charAt(i)) {
        return -1;
      }
    }
    return 0;
  }
}

Solution 4

You should create a unique comparator for comparing the keys of the map. But because you want to print the values too, you should compare the whole entrysets instead:

Comparator<Map.Entry<String, String>> c = new Comparator<Map.Entry<String, String>>() {
  @Override
  public int compare(Map.Entry<String, String> o1, Map.Entry<String, String> o2) {
    int q = Integer.compare(o1.getKey().length(), o2.getKey().length());
    return q != 0 ? q : o1.getKey().compareTo(o2.getKey());
  }
};

Then you can use this comparator in sorting:

map.entrySet().stream().sorted(c).forEach(System.out::println);

View more solutions

Comments

[ernest_k]

Please include what you've done so far to implement this

[Arun Sudhakaran]

whats your criteria for sorting

[Eran]

Since you haven't provided any code and tagged this as Java, I suggest you read the documentation for the Comparator interface and its .thenComparing method

[ZINLIN HTET]

Question>>In short key name order, if the same length, make it alphabetical dictionary order

[Luis Colorado]

please read How to create a Minimal, Complete, and Verifiable example to see why we always ask for the code.

[Andy Turner]

Also: Integer.compare(s1.length(), s2.length()) would be an easier way of writing this compare method (or forgo writing your own class and use comparing(String::length)).

[Darshan Mehta]

@AndyTurner agree. It's another approach I'd say, and upto OP to decide which approach to go ahead with.

[Rogue]

oh hey, someone beat me to posting the Comparator solutions. Way nicer than doing the actual logistics yourself (o1 < o2 barf)