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]