帮前同事关于树结构类型数据获取子节点代码更新

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]