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]