package demo;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.commons.collections.CollectionUtils;
import org.junit.Test;
import java.util.*;
import java.util.stream.Collectors;
public class TreeTest {
/**
* 获取所有节点所有子节点列表
*
* @throws Exception 当json序列化失败时抛出异常
*/
@Test
public void test1() throws Exception {
//组装数据
TreeNode node1 = new TreeNode(1, 0, "根节点1");
TreeNode node2 = new TreeNode(2, 1, "二级节点1");
TreeNode node3 = new TreeNode(3, 1, "二级节点2");
TreeNode node4 = new TreeNode(4, 2, "三级节点1");
TreeNode node5 = new TreeNode(5, 3, "三级节点2");
TreeNode node6 = new TreeNode(6, 3, "根节点2");
TreeNode node7 = new TreeNode(7, 5, "二级节点1");
TreeNode node8 = new TreeNode(8, 6, "二级节点2");
TreeNode node9 = new TreeNode(9, 7, "三级节点1");
TreeNode node10 = new TreeNode(10, 8, "三级节点2");
List<TreeNode> nodes = Arrays.asList(node1,node2,node3,node4,node5,node6,node7,node8,node9,node10);
//将list转为map,以pid为key组装子节点集合
Map<Integer, Set<TreeNode>> childrenNodeMap = nodes.stream().collect(Collectors.groupingBy(TreeNode::getPid, Collectors.toSet()));
//将list转为map，存储key为id,value为node的map
Map<Integer, TreeNode> allNodeMap = nodes.stream().collect(Collectors.toMap(TreeNode::getId, node -> node));
//获取没有子节点的节点信息放入队列
Set<Integer> keySet = childrenNodeMap.keySet();
Deque<Integer> nodeDeque = nodes.stream().map(TreeNode::getId)
.filter(o -> !keySet.contains(o)).collect(Collectors.toCollection(ArrayDeque::new));
Integer currentNodeId;
while ((currentNodeId = nodeDeque.pollFirst()) != null) {
//获取当前节点
TreeNode node = allNodeMap.get(currentNodeId);
if (node != null) {
//获取父节点ID
int pid = node.getPid();
//将当前节点得所有子节点赋值给父节点
Set<TreeNode> nodeList = childrenNodeMap.get(currentNodeId);
if (CollectionUtils.isNotEmpty(nodeList)) {
//使用set去重
childrenNodeMap.get(pid).addAll(nodeList);
}
//父节点入列
nodeDeque.addLast(pid);
}
}
//测试
ObjectMapper objectMapper = new ObjectMapper();
String s = objectMapper.writeValueAsString(childrenNodeMap);
System.out.println(s);
}
/**
* 通过某个节点获取其所有子节点
*/
@Test
public void test2() {
//组装数据
TreeNode node1 = new TreeNode(1, 0, "根节点1");
TreeNode node2 = new TreeNode(2, 1, "二级节点1");
TreeNode node3 = new TreeNode(3, 1, "二级节点2");
TreeNode node4 = new TreeNode(4, 2, "三级节点1");
TreeNode node5 = new TreeNode(5, 3, "三级节点2");
TreeNode node6 = new TreeNode(6, 3, "根节点2");
TreeNode node7 = new TreeNode(7, 5, "二级节点1");
TreeNode node8 = new TreeNode(8, 6, "二级节点2");
TreeNode node9 = new TreeNode(9, 7, "三级节点1");
TreeNode node10 = new TreeNode(10, 8, "三级节点2");
List<TreeNode> nodes = Arrays.asList(node1,node2,node3,node4,node5,node6,node7,node8,node9,node10);
//将list转为map,以pid为key组装子节点集合
Map<Integer, Set<TreeNode>> childrenNodeMap = nodes.stream().collect(Collectors.groupingBy(TreeNode::getPid, Collectors.toSet()));
//当前节点id
Deque<Integer> nodeDeque = new ArrayDeque<>();
//将待搜索节点ID放入队列
nodeDeque.add(3);
Integer currentNodeId;
Set<Integer> resultList = new HashSet<>();
while ((currentNodeId = nodeDeque.pollFirst()) != null) {
resultList.add(currentNodeId);
Set<TreeNode> treeNodes = childrenNodeMap.get(currentNodeId);
if (CollectionUtils.isNotEmpty(treeNodes)) {
Set<Integer> childIdSet = treeNodes.stream().map(TreeNode::getId).collect(Collectors.toSet());
resultList.addAll(childIdSet);
nodeDeque.addAll(childIdSet);
}
}
System.out.println(resultList);
}
}
/**
* 节点类
*
* @author Wangyangyang
*/
@JsonInclude(JsonInclude.Include.NON_EMPTY)
class TreeNode {
/**
* 当前节点id
*/
private Integer id;
/**
* 父节点id
*/
private Integer pid;
/**
* 节点名称
*/
private String name;
public TreeNode(int id, int pid, String name) {
this.id = id;
this.pid = pid;
this.name = name;
}
public TreeNode() {
}
public Integer getId() {
return id;
}
public void setId(Integer id) {
this.id = id;
}
public Integer getPid() {
return pid;
}
public void setPid(Integer pid) {
this.pid = pid;
}
public String getName() {
return name;
}
public void setName(String name) {
this.name = name;
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
TreeNode treeNode = (TreeNode) o;
return Objects.equals(id, treeNode.id) && Objects.equals(pid, treeNode.pid) && Objects.equals(name, treeNode.name);
}
@Override
public int hashCode() {
return Objects.hash(id, pid, name);
}
@Override
public String toString() {
return "Node{" +
"id=" + id +
", pid=" + pid +
'}';
}
}
test1测试结果：
{
"0": [
{
"id": 2,
"pid": 1,
"name": "二级节点1"
},
{
"id": 7,
"pid": 5,
"name": "二级节点1"
},
{
"id": 1,
"pid": 0,
"name": "根节点1"
},
{
"id": 3,
"pid": 1,
"name": "二级节点2"
},
{
"id": 8,
"pid": 6,
"name": "二级节点2"
},
{
"id": 5,
"pid": 3,
"name": "三级节点2"
},
{
"id": 6,
"pid": 3,
"name": "根节点2"
},
{
"id": 10,
"pid": 8,
"name": "三级节点2"
},
{
"id": 4,
"pid": 2,
"name": "三级节点1"
},
{
"id": 9,
"pid": 7,
"name": "三级节点1"
}
],
"1": [
{
"id": 2,
"pid": 1,
"name": "二级节点1"
},
{
"id": 7,
"pid": 5,
"name": "二级节点1"
},
{
"id": 3,
"pid": 1,
"name": "二级节点2"
},
{
"id": 8,
"pid": 6,
"name": "二级节点2"
},
{
"id": 5,
"pid": 3,
"name": "三级节点2"
},
{
"id": 6,
"pid": 3,
"name": "根节点2"
},
{
"id": 10,
"pid": 8,
"name": "三级节点2"
},
{
"id": 4,
"pid": 2,
"name": "三级节点1"
},
{
"id": 9,
"pid": 7,
"name": "三级节点1"
}
],
"2": [
{
"id": 4,
"pid": 2,
"name": "三级节点1"
}
],
"3": [
{
"id": 7,
"pid": 5,
"name": "二级节点1"
},
{
"id": 8,
"pid": 6,
"name": "二级节点2"
},
{
"id": 5,
"pid": 3,
"name": "三级节点2"
},
{
"id": 6,
"pid": 3,
"name": "根节点2"
},
{
"id": 10,
"pid": 8,
"name": "三级节点2"
},
{
"id": 9,
"pid": 7,
"name": "三级节点1"
}
],
"5": [
{
"id": 7,
"pid": 5,
"name": "二级节点1"
},
{
"id": 9,
"pid": 7,
"name": "三级节点1"
}
],
"6": [
{
"id": 8,
"pid": 6,
"name": "二级节点2"
},
{
"id": 10,
"pid": 8,
"name": "三级节点2"
}
],
"7": [
{
"id": 9,
"pid": 7,
"name": "三级节点1"
}
],
"8": [
{
"id": 10,
"pid": 8,
"name": "三级节点2"
}
]
}
test2测试结果：
[3, 5, 6, 7, 8, 9, 10]
