TreeUtils支持同一层级兄弟节点的排序

commons/common-util/src/main/java/com/schbrain/common/util/TreeUtils.java

@@ -2,6 +2,7 @@ package com.schbrain.common.util;
 
 import cn.hutool.core.collection.ListUtil;
 import com.google.common.collect.Maps;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.collections4.MapUtils;
 
@@ -16,6 +17,7 @@ import java.util.function.Function;
  * @author panwangnan
  * @since 2023/7/29
  */
+@Slf4j
 public class TreeUtils {
 
     public static <T, K, C extends Collection<T>> List<T> buildTree(C nodes,
@@ -45,13 +47,38 @@ public class TreeUtils {
         if (CollectionUtils.isEmpty(nodes)) {
             return new ArrayList<>();
         }
+        Map<K, List<T>> parentWithSubNodes = buildParentWithSubNodes(nodes, parentKeyExtractor, parentId);
+        return doBuildTree(keyExtractor, childrenSetter, childMapper, parentWithSubNodes, childrenComparator, parentId);
+    }
+
+    private static <T, K, C extends Collection<T>> Map<K, List<T>> buildParentWithSubNodes(C nodes,
+                                                                                           Function<T, K> parentKeyExtractor,
+                                                                                           @Nullable K parentId) {
         Map<K, List<T>> parentWithSubNodes = StreamUtils.groupBy(nodes, parentKeyExtractor, true);
         if (parentId == null) {
             // groupBy 不允许 key 为空，当 parentId 为空时，单独处理下
             List<T> subNodes = StreamUtils.filterToList(nodes, node -> parentKeyExtractor.apply(node) == null);
             parentWithSubNodes.put(null, subNodes);
         }
-        return doBuildTree(keyExtractor, childrenSetter, childMapper, parentWithSubNodes, childrenComparator, parentId);
+        return parentWithSubNodes;
+    }
+
+    /**
+     * @param preBrotherFlag true: previousId false:nextId
+     */
+    public static <T, K, C extends Collection<T>> List<T> buildTree(C nodes,
+                                                                    Function<T, K> keyExtractor,
+                                                                    Function<T, K> parentKeyExtractor,
+                                                                    Function<T, K> brotherKeyExtractor,
+                                                                    BiConsumer<T, List<T>> childrenSetter,
+                                                                    @Nullable K parentId, @Nullable K rootBrotherId,
+                                                                    boolean preBrotherFlag) {
+        if (CollectionUtils.isEmpty(nodes)) {
+            return new ArrayList<>();
+        }
+        Map<K, List<T>> parentWithSubNodes = buildParentWithSubNodes(nodes, parentKeyExtractor, parentId);
+        return doBuildTree(keyExtractor, brotherKeyExtractor, childrenSetter, parentWithSubNodes, parentId,
+                rootBrotherId, preBrotherFlag);
     }
 
     public static <T, E> List<E> getParents(E self, Collection<T> nodes, Function<T, E> keyMapper, Function<T, E> parentMapper, boolean includeSelf) {
@@ -123,4 +150,52 @@ public class TreeUtils {
         }
     }
 
+    private static <K, T> List<T> doBuildTree(Function<T, K> keyExtractor,
+                                              Function<T, K> brotherKeyExtractor,
+                                              BiConsumer<T, List<T>> childrenSetter,
+                                              Map<K, List<T>> parentWithSubNodes,
+                                              K parentId, K rootBrotherId, boolean preBrotherFlag) {
+        List<T> subNodes = parentWithSubNodes.remove(parentId);
+        if (CollectionUtils.isEmpty(subNodes)) {
+            return Collections.emptyList();
+        }
+        subNodes.forEach(subNode -> {
+            List<T> children = doBuildTree(keyExtractor, brotherKeyExtractor, childrenSetter,
+                    parentWithSubNodes, keyExtractor.apply(subNode), rootBrotherId, preBrotherFlag);
+            childrenSetter.accept(subNode, children);
+        });
+        sort(subNodes, keyExtractor, brotherKeyExtractor, rootBrotherId, preBrotherFlag);
+        return subNodes;
+    }
+
+    private static <K, T> void sort(List<T> dataList, Function<T, K> keyExtractor, Function<T, K> brotherKeyExtractor,
+                                    K rootBrotherId, boolean preBrotherFlag) {
+        if (CollectionUtils.isEmpty(dataList)) {
+            return;
+        }
+        List<T> sortList = new ArrayList<>(dataList.size());
+        Map<K, T> brother2DataMap = StreamUtils.toMap(dataList, brotherKeyExtractor);
+        K currentId = rootBrotherId;
+        while (true) {
+            T current = brother2DataMap.remove(currentId);
+            if (null != current) {
+                currentId = keyExtractor.apply(current);
+                sortList.add(current);
+            } else {
+                break;
+            }
+        }
+        //避免脏数据，则直接返回原列表
+        if (sortList.size() != dataList.size()) {
+            System.out.println(JacksonUtils.toJsonString(sortList));
+            log.warn("Sort Warning:{}", JacksonUtils.toJsonString(dataList));
+            return ;
+        }
+        if (!preBrotherFlag) {
+            Collections.reverse(sortList);
+        }
+        dataList.clear();
+        dataList.addAll(sortList);
+    }
+
 }