TreeMap with comparator comparing length of string in keys

user_1444925 0
Viewed 160

I need a TreeMap mapping Strings to Integers using the length of the string as the sorting criteria. This is my code:

TreeMap<String, Integer> map = new TreeMap<>(Comparator.comparingInt(String::length)) {{
  put("one", 1);
  put("two", 2);
  put("three", 3);
  put("four", 4);
  put("five", 5);
  put("six", 6);
}};

I expect the content of the map to be something like this:

{one=1, two=2, six=6, four=4, five=5, three=3}

but what I get instead is:

{one=6, four=5, three=3}

Looking at the code of the put method of the TreeMap, I can see that if a Comparator has been defined, then it uses the value returned by the comparator to decide if to create a new key in the map, or to overwrite the value of an existing one. In my case, since I have three keys of length 3 ("one", "two" and "six"), it inserts only one key of length 3 (the first inserted, "one") and updates its value (first 1, then 2 and at last 6). The same goes for key of length 4 ("four" and "five").

How can I have the tree map inserting all the keys I defined ("one", "two", "three", "four", "five", "six") sorted by the length of the key?

treemap java
2 Answers

The specified order {one=1, two=2, six=6, four=4, five=5, three=3} cannot be achieved with the TreeMap because this implementation can be sorted only by keys, and as soon as the lengths of the keys have multiple collisions, additional comparator for the value is needed.

Thus, a raw map may be created first, then its entrySet can be sorted in a custom order using Stream API Stream::sorted and Collectors.toMap:

  1. by key length (ascending)
  2. by value (ascending)

Then the result is collected into a LinkedHashMap famous for maintaining the insertion order:

Map<String, Integer> raw = new HashMap<>() {{
  put("one", 1);
  put("two", 2);
  put("three", 3);
  put("four", 4);
  put("five", 5);
  put("six", 6);
}};

Map<String, Integer> map = raw.entrySet()
    .stream()
    .sorted(Comparator.<Map.Entry<String, Integer>>comparingInt(e -> e.getKey().length())
        .thenComparingInt(Map.Entry::getValue))
    .collect(Collectors.toMap(
        Map.Entry::getKey,
        Map.Entry::getValue,
        (a, b) -> a,
        LinkedHashMap::new
    ));
System.out.println(map);

Output:

{one=1, two=2, six=6, four=4, five=5, three=3}

Update

It appears that for the given input a custom comparator may be provided for the TreeMap maintaining the insertion order:

Map<String, Integer> map2 = new TreeMap<>(
    Comparator.comparingInt(String::length)
        .thenComparing((k1, k2) -> 1) // pick the "earlier" key of the two
) {{
  put("one", 1);
  put("two", 2);
  put("three", 3);
  put("four", 4);
  put("five", 5);
  put("six", 6);
}};
System.out.println(map2);

Output:

{one=1, two=2, six=6, four=4, five=5, three=3}
user_132798310

You need a double comparator, where you first compare the string lengths, then the strings themselves. Like:

public class MyTreeMapComparator implements Comparator<String> {
  public int compare(String one, String two) {
    if (one.length() == two.length()) {
      return one.compareTo(two);
    } else return two.length() - one.length();
  }
}
user_12914920
Related