在b站上刷到了一个实现类似功能的视频,感觉不是特别灵活。故也自实现了一版。
代码
import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.Deque;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.function.Predicate;
/**
* Created by Kim QQ.Cong on 2022-04-24 / 21:18
*
* @author: CongQingquan
* @Description: 通用树工具类
*/
public class TreeUtils {
private TreeUtils() {
}
/**
* 树化方法的参数封装接口
*
* @param <T>
* @param <ID>
*/
public interface ToTreeParam<T, ID> {
/**
* 获取id提取器 (id & pid同类型)
*/
Function<T, ID> getIdExtractor();
/**
* 获取pid提取器 (id & pid同类型)
*/
Function<T, ID> getPidExtractor();
/**
* 获取节点排序比较器
*/
Comparator<? super T> getComparator();
/**
* 获取设置子节点的BiConsumer
*/
BiConsumer<T, Collection<T>> getSettingChildrenBiConsumer();
/**
* 获取默认的根元素断言 (test: pid == null)
*/
default Predicate<T> getRootPredicate() {
return (n) -> this.getPidExtractor().apply(n) == null;
}
}
/**
* 树化
*
* @param param 参数封装类
* @param <T>
* @param <ID>
* @return
*/
public static <T, ID> Collection<T> toTree(Collection<T> nodes, ToTreeParam<T, ID> param) {
if (param == null) {
throw new IllegalArgumentException("To tree param cannot be null");
}
return toTree(nodes, param.getIdExtractor(), param.getPidExtractor(), param.getRootPredicate(),
param.getComparator(), param.getSettingChildrenBiConsumer());
}
/**
* 树化
* 1. id data type === parent id data type
* 2. 默认的,root node: pidExtractor.apply(node) == null & child node: pidExtractor.apply(node) != null
*
* @param nodes 节点集合
* @param idExtractor id提取器 (id & pid同类型)
* @param pidExtractor pid提取器 (id & pid同类型)
* @param <T> 节点类型
* @param <ID> id数据类型 (id & pid同类型)
* @param isRoot 根节点断言
* @param comparator 节点排序比较器
* @param settingChildren 设置子节点
* 1. 为什么用 set 的方式而不是 get 的方式?
* 考虑 get 的 children collection 可能为 null 的情况,初始化后无法进行 children 数据的写回,希望调用者提前写初始化操作
* 2. 调用者在初始化 children 容器时,根据容器底层使用的数据结构,可能会影响到 children 节点根据 comparator 的排序顺序
*
* @return
*/
public static <T, ID> Collection<T> toTree(final Collection<T> nodes,
final Function<T, ID> idExtractor,
final Function<T, ID> pidExtractor,
final Predicate<T> isRoot,
final Comparator<? super T> comparator,
final BiConsumer<T, Collection<T>> settingChildren) {
if (nodes == null) {
return new ArrayList<>();
}
if (nodes.size() == 0) {
return nodes;
}
if (idExtractor == null || pidExtractor == null) {
throw new IllegalArgumentException("Extractor cannot be null");
}
if (isRoot == null) {
throw new IllegalArgumentException("Root node predicate cannot be null");
}
final Set<T> roots = new TreeSet<>(comparator);
final Map<ID, Collection<T>> childrenMapping = new HashMap<>(16);
nodes.forEach(node -> {
if (isRoot.test(node)) {
roots.add(node);
return;
}
childrenMapping.computeIfAbsent(pidExtractor.apply(node), (id) -> new TreeSet<>(comparator)).add(node);
});
final Deque<T> stack = new LinkedList<>();
roots.forEach(root -> {
stack.add(root);
while (stack.size() > 0) {
T currentNode = stack.pop();
ID currentNodeId = idExtractor.apply(currentNode);
Collection<T> cc = childrenMapping.get(currentNodeId);
if (CollectionUtils.isEmpty(cc)) {
continue;
}
for (T cn : cc) {
stack.push(cn);
}
settingChildren.accept(currentNode, cc);
}
});
return roots;
}
}
代码设计
利用Stack而非递归的形式实现的子节点搜索, 并且节点类不需要继承任何类似Node这样的内部提供的Base类。 一般来说可能需要通过定义一个抽象的Node类,里面定义好最基本的字段,如:id、pid、children。 并提供对应的set/get方法。但是这样需要强制节点类继承Node类,如果实体类已有继承类呢?如果实体类已经有了其中的某些字段呢?约束过多,体验不好。所以我想通过lambda传参的形式,在配合泛型,来完成这些操作。
比如:
- id & pid 的提取器,获取实体类的id & pid字段值。这里的id&pid是抽象的,具体是实体类的哪个字段提供,由调用者决定。
- isRoot 的断言,这是刚才没有提到的。如何判断节点为一级根节点,不能仅仅通过常规的pid==null的形式来判定,这无疑又进行了一次强制约束。比如,我就是想要将二级节点作为根节点,且节点的pid就是不为空。所以,还是交给调用者来决定。
- settingChildren,为什么设定children也被提取出来?有两层含义:
- 调用者来确定children对应实体类的哪个字段
- 不确定获取的实体类的children集合字段是否为null,如果为null,那么无法进行children数据的写回
参数简化
动态的代价基本都是实现代码的增加 或 调用代码增加。如:
Collection<Node> nodes = TreeUtils.toTree(roots, Node::getId, Node::getParentId,
(node) -> node.getParentId() == null,
Comparator.comparing(Node::getId).reversed(),
(node, cns) -> {
List<Node> cl = node.getChildren();
if (cl == null) {
node.setChildren(Lists.newArrayList(cns));
return;
}
cl.addAll(cns);
});
那么为了进行简化,内部提供了参数封装接口:ToTreeParam。简化后的代码如下:
class DefaultToTreeParam implements ToTreeParam<Node, String> {
@Override
public Function<Node, String> getIdExtractor() {
return Node::getId;
}
@Override
public Function<Node, String> getPidExtractor() {
return Node::getParentId;
}
@Override
public Predicate<Node> getRootPredicate() {
return (n) -> n.getId().length() == 2;
}
@Override
public Comparator<? super Node> getComparator() {
return Comparator.comparing(Node::getId).reversed();
}
@Override
public BiConsumer<Node, Collection<Node>> getSettingChildrenBiConsumer() {
return (node, cns) -> {
List<Node> cl = node.getChildren();
if (cl == null) {
node.setChildren(Lists.newArrayList(cns));
return;
}
cl.addAll(cns);
};
}
}
Collection<Node> nodes2 = TreeUtils.toTree(data, new DefaultToTreeParam());
嗯?代码怎么多了?因为这仅是一次调用。当需要在不同的多个类中调用TreeUtils.toTree时,这种方式更为简洁方便。