实现一个通用的树形结构构建工具

文章目录

1. 前言
2. 树结构
3. 具体实现逻辑
- 3.1 TreeNode
- 3.2 TreeUtils
- 3.3 例子
4. 小结

1. 前言

树结构的生成在项目中应该都比较常见，比如部门结构树的生成，目录结构树的生成，但是大家有没有想过，如果在一个项目中有多个树结构，那么每一个都要定义一个生成方法显然是比较麻烦的，所以我们就想写一个通用的生成树方法，下面就来看下如何来写。

2. 树结构

在这里插入图片描述
看上面的图，每一个节点都会有三个属性

parentId 表示父节点 ID，根节点的父结点 ID = null
id 表示当前节点 ID，这个 ID 用来标识一个节点
children 是当前节点的子节点

那么上面来介绍完基本的几个属性，下面就来看下具体的实现了。

3. 具体实现逻辑

3.1 TreeNode

TreeNode 是公共节点，就是顶层父类，里面的属性就是上面图中的三个。

@Data
@AllArgsConstructor
@NoArgsConstructor
@Accessors(chain = true)
public class TreeNode<T, V> {private T parentId;private T id;private List<TreeNode<T, V>> children;public TreeNode(T parentId, T id) {this.parentId = parentId;this.id = id;}public void addChild(TreeNode<T, V> treeNode){if(children == null){children = new ArrayList<>();}children.add(treeNode);}}

TreeNode 里面的 id 都是用的范型，其中 T 就是 id 的类型，因为这个 id 有可能是 Long、Int、String … 类型，不一定是 Long。另一个 V 就是具体的节点类型。

使用范型的好处就是扩展性高，不需要把属性写死。

3.2 TreeUtils

这个是工具类，专门实现树的构建以及一些其他的方法，下面一个一个来看。首先是创建树的方法：

/*** 构建一棵树** @param flatList* @param <T>* @param <V>* @return*/
public static <T, V extends TreeNode<T, V>> List<V> buildTree(List<V> flatList) {if (flatList == null || flatList.isEmpty()) {return null;}Map<T, TreeNode<T, V>> nodeMap = new HashMap<>();for (TreeNode<T, V> node : flatList) {nodeMap.put(node.getId(), node);}// 查找根节点List<V> rootList = new ArrayList<>();for (V node : flatList) {// 如果父节点为空，就是一个根节点if (node.getParentId() == null) {rootList.add(node);} else {// 父节点不为空，就是子节点TreeNode<T, V> parent = nodeMap.get(node.getParentId());if (parent != null) {parent.addChild(node);} else {rootList.add(node);}}}return rootList;
}

整体时间复杂度：O(n)，创建的时候传入节点集合，然后返回根节点集合。里面的逻辑是首先放到一个 nodeMap 中，然后遍历传入的集合，根据 parentId 进行不同的处理。逻辑不难，看注释即可。但是创建树的时候，有时候我们希望根据某个顺序对树进行排序，比如同一层的我想根据名字或者 id 进行排序，顺序或者倒序都可以，那么就可以使用下面的方法。

/**
* 构建一棵排序树
*
* @param flatList
* @param comparator
* @param <T>
* @param <V>
* @return
*/
public static <T, V extends TreeNode<T, V>> List<V> buildTreeWithCompare(List<V> flatList, Comparator<V> comparator) {if (flatList == null || flatList.isEmpty()) {return Collections.emptyList(); // 返回空列表而不是null，这通常是一个更好的实践}// 子节点分组Map<T, List<V>> childGroup = flatList.stream().filter(v -> v.getParentId() != null).collect(Collectors.groupingBy(V::getParentId));// 找出父节点List<V> roots = flatList.stream().filter(v -> v.getParentId() == null).sorted(comparator) // 根据提供的比较器对根节点进行排序.collect(Collectors.toList());// 构建树for (V root : roots) {buildTreeRecursive(root, childGroup, comparator);}return roots;
}private static <T, V extends TreeNode<T, V>> void buildTreeRecursive(V parent, Map<T, List<V>> childGroup, Comparator<V> comparator) {List<V> children = childGroup.get(parent.getId());if (children != null) {// 对子节点进行排序children.sort(comparator);// 将排序后的子节点添加到父节点中children.forEach(parent::addChild);// 递归对子节点继续处理children.forEach(child -> buildTreeRecursive(child, childGroup, comparator));}
}

这里面是使用的递归，其实也可以使用层次遍历的方式来写，或者直接用第一个 buildTree 方法来往里面套也行。

上面这两个是关键的方法，那么下面再给出一些其他的非必要方法，比如查询节点数。下面这个方法就是获取以 root 为根的数的节点数。

/*** 查询以 root 为根的树的节点数** @param root* @param <T>* @param <V>* @return*/
private static <T, V extends TreeNode<T, V>> long findTreeNodeCount(TreeNode<T, V> root) {if (root == null) {return 0;}long res = 1;List<TreeNode<T, V>> children = root.getChildren();if (children == null || children.isEmpty()) {return res;}for (TreeNode<T, V> child : children) {res += findTreeNodeCount(child);}return res;
}

上面是传入一个根节点，获取这棵树的节点数，而下面的就是传入一个集合来分别获取节点数，里面也是调用了上面的 findTreeNodeCount 方法去获取。

/*** 查询给定集合的节点数** @param nodes 根节点集合* @param <T>* @param <V>* @return*/
public static <T, V extends TreeNode<T, V>> HashMap<V, Long> findTreeNodeCount(List<V> nodes) {if (nodes == null || nodes.isEmpty()) {return new HashMap<>(); // 返回空列表而不是null，这通常是一个更好的实践}HashMap<V, Long> map = new HashMap<>();for (V root : nodes) {map.put(root,  findTreeNodeCount(root));}return map;
}

下面再给一下获取数的深度的方法。

// 查找树的最大深度
private static <T, V extends TreeNode<T, V>> int getMaxDepthV(TreeNode<T, V> root) {if (root == null || root.getChildren() == null || root.getChildren().isEmpty()) {return 1;}return 1 + root.getChildren().stream().mapToInt(TreeUtils::getMaxDepthV).max().getAsInt();
}public static <T, V extends TreeNode<T, V>> int getMaxDepth(V root) {return getMaxDepthV(root);
}

最后，我们拿到一棵树之后，肯定有时候会希望在里面查找一些具有特定属性的节点，比如某个节点名字是不是以 xx 开头 … ，这时候就可以用下面的方法。

// 查找所有具有特定属性的节点
public static <T, V extends TreeNode<T, V>> List<V> findAllNodesByProperty(TreeNode<T, V> root, Function<V, Boolean> predicate) {if (root == null) {return Collections.emptyList();}List<V> result = new ArrayList<>();// 符合属性值if (predicate.apply((V) root)) {result.add((V) root);}if (root.getChildren() == null || root.getChildren().isEmpty()) {return result;}for (TreeNode<T, V> child : root.getChildren()) {result.addAll(findAllNodesByProperty(child, predicate));}return result;
}

好了，方法就这么多了，其他方法如果你感兴趣也可以继续补充下去，那么这些方法是怎么用的呢？范型的好处要怎么体现呢？下面就来看个例子。

3.3 例子

首先我们有一个部门类，里面包括部门的名字，然后我需要对这个部门集合来构建一棵部门树。

@Data
@ToString
@NoArgsConstructor
public class Department extends TreeNode<String, Department> {private String name;public Department(String id, String parentId, String name){super(parentId, id);this.name = name;}}

构建的方法如下：

public class Main {public static void main(String[] args) {List<Department> flatList = new ArrayList<>();flatList.add(new Department("1", null, "Sales"));flatList.add( new Department("2", "1", "East Sales"));flatList.add( new Department("3", "1","West Sales"));flatList.add( new Department("4", "2","East Sales Team 1"));flatList.add( new Department("5", "2","East Sales Team 2"));flatList.add( new Department("6", "3","West Sales Team 1"));List<Department> departments = TreeUtils.buildTreeWithCompare(flatList, (o1, o2) -> {return o2.getName().compareTo(o1.getName());});Department root = departments.get(0);List<Department> nodes = TreeUtils.findAllNodesByProperty(root, department -> department.getName().startsWith("East"));System.out.println(nodes);System.out.println(TreeUtils.getMaxDepth(root));System.out.println(TreeUtils.findTreeNodeCount(nodes));}}

可以看下 buildTreeWithCompare 的输出：
在这里插入图片描述
其他的输出如下：

[Department(name=East Sales), Department(name=East Sales Team 2), Department(name=East Sales Team 1)]
3
{Department(name=East Sales)=3, Department(name=East Sales Team 2)=1, Department(name=East Sales Team 1)=1}