TreeMap自定义排序

我们都知道TreeMap可以根据key按字典升序排序。但在某些场景下，我们需要自定义排序规则，为了代码优雅一些，我们也希望在stream中groupingBy时自定义排序规则，就可以参考本文的实现。

1. 使用TreeMap默认的排序规则（按字典升序排序）

先来看一下默认的排序规则。

/*** 使用TreeMap默认的排序规则（按字典升序排序）*/
private static void sortByDefault() {TreeMap<String, Integer> treeMap = new TreeMap<>();for (int index = 0; index < 5; index++) {treeMap.put("str_" + index, index);}System.out.println("*****TreeMap默认排序规则*****");treeMap.forEach((k, v) -> System.out.println(k + " : " + v));System.out.println();
}

2. 使用自定义排序规则

使用自定义排序规则对TreeMap进行实例化。

/*** 使用自定义排序规则*/
private static void sortByCustom() {Comparator<String> comparator = (String s1, String s2) -> SortEnum.getSortWeightByValue(s1) - SortEnum.getSortWeightByValue(s2);TreeMap<String, Integer> treeMap = new TreeMap<>(comparator);for (int index = 0; index < 5; index++) {treeMap.put("str_" + index, index);}System.out.println("*****TreeMap自定义排序规则*****");treeMap.forEach((k, v) -> System.out.println(k + " : " + v));System.out.println();
}

3. GroupingBy时使用自定义排序规则

一开始直接用TreeMap::new作为groupingBy方法的第二个参数，但是这种办法只能引用TreeMap的无参构造方法。查看groupingBy的入参可以发现，第二个参数是Supplier类型，于是手动实现Supplier，在方法体中传入自定义比较器，然后返回TreeMap对象。

/*** GroupingBy时使用自定义排序规则*/
private static void sortWhenGroupingBy() {List<String> list = new ArrayList<>();for (int index = 0; index < 5; index++) {list.add("str_" + index);}Comparator<String> comparator = (s1, s2) -> SortEnum.getSortWeightByValue(s1) - SortEnum.getSortWeightByValue(s2);// 如果现在外面定义好Supplier，需要声明TreeMap的key和value的类型，否则可能出现类型不兼容的异常// Supplier<TreeMap<String, List<String>>> supplier = () -> new TreeMap<>(comparator);// TreeMap<String, List<String>> treeMap = list.stream().collect(Collectors.groupingBy(Function.identity(), supplier, Collectors.toList()));TreeMap<String, List<String>> treeMap = list.stream().collect(Collectors.groupingBy(Function.identity(), () -> new TreeMap<>(comparator), Collectors.toList()));System.out.println("*****GroupingBy时自定义排序规则*****");treeMap.forEach((k, v) -> System.out.println(k + " : " + v));System.out.println();
}

完整代码

自定义排序权重类

public enum SortEnum {ZERO(0, "str_1"),FIRST(1, "str_0"),SECOND(2, "str_3"),THREE(3, "str_2"),;private int sortWeight;private String value;SortEnum(int sortWeight, String value) {this.sortWeight = sortWeight;this.value = value;}public static int getSortWeightByValue(String value) {for (SortEnum sortEnum : SortEnum.values()) {if (sortEnum.value.equals(value)) {return sortEnum.sortWeight;}}return Integer.MAX_VALUE;}
}

测试代码

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.TreeMap;
import java.util.function.Function;
import java.util.stream.Collectors;public class TestTreeMap {public static void main(String[] args) {sortByDefault();sortByCustom();sortWhenGroupingBy();}/*** 使用TreeMap默认的排序规则（按字典升序排序）*/private static void sortByDefault() {TreeMap<String, Integer> treeMap = new TreeMap<>();for (int index = 0; index < 5; index++) {treeMap.put("str_" + index, index);}System.out.println("*****TreeMap默认排序规则*****");treeMap.forEach((k, v) -> System.out.println(k + " : " + v));System.out.println();}/*** 使用自定义排序规则*/private static void sortByCustom() {Comparator<String> comparator = (String s1, String s2) -> SortEnum.getSortWeightByValue(s1) - SortEnum.getSortWeightByValue(s2);TreeMap<String, Integer> treeMap = new TreeMap<>(comparator);for (int index = 0; index < 5; index++) {treeMap.put("str_" + index, index);}System.out.println("*****TreeMap自定义排序规则*****");treeMap.forEach((k, v) -> System.out.println(k + " : " + v));System.out.println();}/*** GroupingBy时使用自定义排序规则*/private static void sortWhenGroupingBy() {List<String> list = new ArrayList<>();for (int index = 0; index < 5; index++) {list.add("str_" + index);}Comparator<String> comparator = (s1, s2) -> SortEnum.getSortWeightByValue(s1) - SortEnum.getSortWeightByValue(s2);// 如果现在外面定义好Supplier，需要声明TreeMap的key和value的类型，否则可能出现类型不兼容的异常// Supplier<TreeMap<String, List<String>>> supplier = () -> new TreeMap<>(comparator);// TreeMap<String, List<String>> treeMap = list.stream().collect(Collectors.groupingBy(Function.identity(), supplier, Collectors.toList()));TreeMap<String, List<String>> treeMap = list.stream().collect(Collectors.groupingBy(Function.identity(), () -> new TreeMap<>(comparator), Collectors.toList()));System.out.println("*****GroupingBy时自定义排序规则*****");treeMap.forEach((k, v) -> System.out.println(k + " : " + v));System.out.println();}
}

运行结果：

*****TreeMap默认排序规则*****
str_0 : 0
str_1 : 1
str_2 : 2
str_3 : 3
str_4 : 4*****TreeMap自定义排序规则*****
str_1 : 1
str_0 : 0
str_3 : 3
str_2 : 2
str_4 : 4*****GroupingBy时自定义排序规则*****
str_1 : [str_1]
str_0 : [str_0]
str_3 : [str_3]
str_2 : [str_2]
str_4 : [str_4]