Search code examples
javatreemap

TreeMap with comparator comparing length of string in keys


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?


Solution

  • 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}