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 0dccc5d189c4949907aec48576d59055849dae14..ea853e2faa7561d29c70be71e5660e16c8cd0fd9 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 @@ -4,6 +4,7 @@ import org.apache.commons.collections4.CollectionUtils; import javax.annotation.Nullable; import java.util.ArrayList; +import java.util.Comparator; import java.util.List; import java.util.Map; import java.util.function.BiConsumer; @@ -31,6 +32,16 @@ public class TreeUtils { Function childMapper, BiConsumer> childrenSetter, @Nullable K parentId) { + return buildTree(nodes, keyExtractor, parentKeyExtractor, childMapper, childrenSetter, null, parentId); + } + + public static List buildTree(List nodes, + Function keyExtractor, + Function parentKeyExtractor, + Function childMapper, + BiConsumer> childrenSetter, + Comparator childrenComparator, + @Nullable K parentId) { if (CollectionUtils.isEmpty(nodes)) { return new ArrayList<>(); } @@ -40,20 +51,24 @@ public class TreeUtils { List subNodes = StreamUtils.filterToList(nodes, node -> parentKeyExtractor.apply(node) == null); parentWithSubNodes.put(null, subNodes); } - return buildTree(keyExtractor, childrenSetter, childMapper, parentWithSubNodes, parentId); + return doBuildTree(keyExtractor, childrenSetter, childMapper, parentWithSubNodes, childrenComparator, parentId); } - private static List buildTree(Function keyExtractor, - BiConsumer> childrenSetter, - Function childMapper, - Map> parentWithSubNodes, - K parentId) { - List subNodes = parentWithSubNodes.get(parentId); + private static List doBuildTree(Function keyExtractor, + BiConsumer> childrenSetter, + Function childMapper, + Map> parentWithSubNodes, + Comparator childrenComparator, + K parentId) { + List subNodes = parentWithSubNodes.remove(parentId); return StreamUtils.toList(subNodes, subNode -> { - E convertedSubNode = childMapper.apply(subNode); - List children = buildTree(keyExtractor, childrenSetter, childMapper, parentWithSubNodes, keyExtractor.apply(subNode)); - childrenSetter.accept(convertedSubNode, children); - return convertedSubNode; + E child = childMapper.apply(subNode); + List children = doBuildTree(keyExtractor, childrenSetter, childMapper, parentWithSubNodes, childrenComparator, keyExtractor.apply(subNode)); + if (CollectionUtils.isNotEmpty(children) && childrenComparator != null) { + children.sort(Comparator.nullsLast(childrenComparator)); + } + childrenSetter.accept(child, children); + return child; }); }