diff --git a/commons/common-util/src/main/java/com/schbrain/common/util/TreeUtils.java b/commons/common-util/src/main/java/com/schbrain/common/util/TreeUtils.java index bf4c9d144e2c7325118b2d28f0bd489fddc4043a..6aa78d25db61a6dc38233a163124d48811a85da5 100644 --- a/commons/common-util/src/main/java/com/schbrain/common/util/TreeUtils.java +++ b/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 > List buildTree(C nodes, @@ -45,13 +47,38 @@ public class TreeUtils { if (CollectionUtils.isEmpty(nodes)) { return new ArrayList<>(); } + Map> parentWithSubNodes = buildParentWithSubNodes(nodes, parentKeyExtractor, parentId); + return doBuildTree(keyExtractor, childrenSetter, childMapper, parentWithSubNodes, childrenComparator, parentId); + } + + private static > Map> buildParentWithSubNodes(C nodes, + Function parentKeyExtractor, + @Nullable K parentId) { Map> parentWithSubNodes = StreamUtils.groupBy(nodes, parentKeyExtractor, true); if (parentId == null) { // groupBy 不允许 key 为空,当 parentId 为空时,单独处理下 List 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 > List buildTree(C nodes, + Function keyExtractor, + Function parentKeyExtractor, + Function brotherKeyExtractor, + BiConsumer> childrenSetter, + @Nullable K parentId, @Nullable K rootBrotherId, + boolean preBrotherFlag) { + if (CollectionUtils.isEmpty(nodes)) { + return new ArrayList<>(); + } + Map> parentWithSubNodes = buildParentWithSubNodes(nodes, parentKeyExtractor, parentId); + return doBuildTree(keyExtractor, brotherKeyExtractor, childrenSetter, parentWithSubNodes, parentId, + rootBrotherId, preBrotherFlag); } public static List getParents(E self, Collection nodes, Function keyMapper, Function parentMapper, boolean includeSelf) { @@ -123,4 +150,52 @@ public class TreeUtils { } } + private static List doBuildTree(Function keyExtractor, + Function brotherKeyExtractor, + BiConsumer> childrenSetter, + Map> parentWithSubNodes, + K parentId, K rootBrotherId, boolean preBrotherFlag) { + List subNodes = parentWithSubNodes.remove(parentId); + if (CollectionUtils.isEmpty(subNodes)) { + return Collections.emptyList(); + } + subNodes.forEach(subNode -> { + List 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 void sort(List dataList, Function keyExtractor, Function brotherKeyExtractor, + K rootBrotherId, boolean preBrotherFlag) { + if (CollectionUtils.isEmpty(dataList)) { + return; + } + List sortList = new ArrayList<>(dataList.size()); + Map 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); + } + }