树结构处理，list和tree互转

1、实体类

package com.iot.common.test.entity;

import lombok.Data;

import java.util.List;

/**
 * @description:
 * @author:zilong
 * @date:2023/9/8
 */
@Data
public class Node {
    //id
    private String id;
    //父节点id
    private String pId;
    //名称
    private String name;
    //编码
    private String code;
    //层级
    private String level;
    private List<Node> children;

    public Node(String id, String pId, String name, String code, String level) {
        this.id = id;
        this.pId = pId;
        this.name = name;
        this.code = code;
        this.level = level;
    }


    public Node() {

    }
}

2、工具类

package com.iot.common.util;

import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.core.toolkit.support.SFunction;
import com.iot.common.test.entity.Node;
import lombok.SneakyThrows;

import java.lang.invoke.SerializedLambda;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树工具类
 */

public class TreeUtil {


    /**
     * 基础数据转树结构Map版本（速度比递归要快很多）
     *
     * @param sourceList  需转换的数据
     * @param getId       主键
     * @param getParentId 父id （父id必须和主键相同类型）
     * @param getChildren 子集
     * @param setChildren 子集
     * @return tree
     */
    public static <T, R> List<T> listToTree(List<T> sourceList, Function<T, R> getId, Function<T, R> getParentId,
                                            Function<T, List<T>> getChildren, BiConsumer<T, List<T>> setChildren) {

        Map<R, T> oldMap = sourceList.stream().collect(Collectors.toMap(getId, T -> T));
        List<T> treeList = new ArrayList<>();
        sourceList.forEach(tree -> {
            T parent = oldMap.get(getParentId.apply(tree));
            if (parent == null) {
                treeList.add(tree);
            } else {
                List<T> ch = getChildren.apply(parent);
                if (ch == null) {
                    ch = new ArrayList<>();
                }
                ch.add(tree);
                setChildren.accept(parent, ch);
            }
        });
        return treeList;
    }

    /**
     * 树结构转父子排序列表且按父在前，子在后，进行排序。
     *
     * @param treeList    父子树列表（带有id、parentId和children的模型列表）
     * @param getId
     * @param getParentId
     * @param getChildren
     * @param setChildren
     * @param <T>
     * @param <R>
     * @return
     */
    public static <T, R> List<T> treeListToSortList(List<T> treeList, Function<T, R> getId, Function<T, R> getParentId,
                                                    Function<T, List<T>> getChildren, BiConsumer<T, List<T>> setChildren) {
        // 先整成树形结构
        List<T> list = listToTree(treeList, getId, getParentId, getChildren, setChildren);
        List<T> sortList = new ArrayList<>();
        tree2List(sortList, list.get(0), getChildren, setChildren);
        return sortList;
    }

    /**
     * 树结构转列表
     *
     * @param result      结果容器
     * @param t           树顶部元素
     * @param getChildren
     * @param <T>
     */
    private static <T> void tree2List(List<T> result, T t, Function<T, List<T>> getChildren, BiConsumer<T, List<T>> setChildren) {
        //根据条件判断是否需要添加至列表
        result.add(t);
        List<T> children = getChildren.apply(t);
        // 将children置成空
        setChildren.accept(t, null);
        //没有子级
        if (children == null || children.size() == 0) {
            return;
        }
        //存在子级，递归调用
        for (T child : children) {
            tree2List(result, child, getChildren, setChildren);
        }
    }

    /**
     * 基础数据集转树结构
     *
     * @param baseList      树结构的基础数据集
     * @param getIdFn       获取主键的函数
     * @param getParentIdFn 获取父节点的函数
     * @param getChildrenFn 获取子集的函数
     * @param <T>           t
     * @param <R>           r
     * @return t
     */
    @SneakyThrows
    public static <T, R> List<T> treeOut(List<T> baseList, Function<T, R> getIdFn, Function<T, R> getParentIdFn, SFunction<T, R> getChildrenFn) {
        /*所有元素的Id*/
        List<Object> ids = baseList.stream().map(getIdFn).collect(Collectors.toList());
        /*查出所有顶级节点*/
        List<T> topLevel = baseList.stream().filter(x -> {
            R apply = getParentIdFn.apply(x);
            return !ids.contains(apply);
        }).collect(Collectors.toList());
        return TreeUtil.recursion(topLevel, baseList, getIdFn, getParentIdFn, getChildrenFn);
    }

    /**
     * 指定顶级元素的基础数据集转树结构
     *
     * @param list
     * @param top
     * @param getIdFn
     * @param getParentIdFn
     * @param getChildrenFn
     * @param <T>
     * @param <R>
     * @return
     */
    @SneakyThrows
    public static <T, R> List<T> treeOutWithTop(List<T> list, T top, Function<T, R> getIdFn, Function<T, R> getParentIdFn, SFunction<T, R> getChildrenFn) {
        ArrayList<T> ts = new ArrayList<>();
        ts.add(top);
        return TreeUtil.recursion(ts, list, getIdFn, getParentIdFn, getChildrenFn);
    }

    @SneakyThrows
    private static <T, R> List<T> recursion(List<T> superLevel, List<T> list, Function<T, R> getIdFn, Function<T, R> getParentIdFn, SFunction<T, R> getChildrenFn) {
        //获取setChildren的Method
        Method writeReplaceMethod = getChildrenFn.getClass().getDeclaredMethod("writeReplace");
        boolean accessible = writeReplaceMethod.isAccessible();
        writeReplaceMethod.setAccessible(true);
        SerializedLambda serializedLambda = (SerializedLambda) writeReplaceMethod.invoke(getChildrenFn);
        writeReplaceMethod.setAccessible(accessible);
        String setMethodName = serializedLambda.getImplMethodName().replaceFirst("g", "s");
        Method setMethod = Class.forName(serializedLambda.getImplClass().replace("/", ".")).getDeclaredMethod(setMethodName, List.class);

        for (T t : superLevel) {
            List<T> children = list.stream().filter(x -> {
                R apply = getParentIdFn.apply(x);
                R apply1 = getIdFn.apply(t);
                return apply.equals(apply1);
            }).collect(Collectors.toList());
            if (children.size() <= 0) {
                continue;
            }

            List<T> recursion = recursion(children, list, getIdFn, getParentIdFn, getChildrenFn);
            setMethod.invoke(t, recursion);
        }
        return superLevel;
    }

    /**
     * 将列表转换为树形结构
     * 注：此方法的根节点的pid需要为null
     *
     * @param nodeList 列表
     * @return 树形结构
     */
    public static List<Node> buildTree(List<Node> nodeList) {
        // 存储所有节点
        Map<String, Node> nodeMap = new HashMap<>(nodeList.size());
        // 存储根节点
        List<Node> rootList = new ArrayList<>();

        // 将所有节点存储到map中，并找出根节点
        for (Node node : nodeList) {
            nodeMap.put(node.getId(), node);
            if (node.getPId() == null) {
                rootList.add(node);
            }
        }

        // 遍历所有节点，将子节点添加到父节点的children中
        for (Node node : nodeList) {
            String pId = node.getPId();
            if (pId != null) {
                Node parentNode = nodeMap.get(pId);
                if (parentNode != null) {
                    if (parentNode.getChildren() == null) {
                        parentNode.setChildren(new ArrayList<>());
                    }
                    parentNode.getChildren().add(node);
                }
            }
        }
        return rootList;
    }

    public static void main(String[] args) {
        ArrayList<Node> list = new ArrayList<Node>() {
            {
                add(new Node("1", "0", "A", "1", "1"));
                add(new Node("2", "0", "B", "2", "1"));
                add(new Node("3", "1", "A-1", "11", "2"));
                add(new Node("4", "1", "A-2", "12", "2"));
                add(new Node("5", "3", "A-1-1", "111", "3"));
                add(new Node("6", "5", "A-1-1-1", "1111", "4"));
                add(new Node("7", "6", "A-1-1-1-1", "11111", "5"));
                add(new Node("8", "2", "B-1", "21", "2"));
                add(new Node("9", "8", "B-1-1", "211", "3"));
                add(new Node("10", "9", "B-1-1-1", "2111", "4"));
            }
        };


        //使用工具类 treeOut:
        List<Node> result = TreeUtil.treeOut(list, Node::getId, Node::getPId, Node::getChildren);
        System.out.println("-----【treeOut】-----");
        System.out.println(JSON.toJSONString(result));

        //使用工具类 listToTree（Map）:
        List<Node> result1 = TreeUtil.listToTree(list, Node::getId, Node::getPId, Node::getChildren, Node::setChildren);
        System.out.println("-----【listToTree】-----");
        System.out.println(JSON.toJSONString(result1));

        //使用工具类 treeOutWithTop:
        Node top = new Node();
        top.setId("3");
        List<Node> result2 = TreeUtil.treeOutWithTop(list, top, Node::getId, Node::getPId, Node::getChildren);
        System.out.println("-----【treeOutWithTop】-----");
        System.out.println(JSON.toJSONString(result2));

        //使用工具类 buildTree:
        List<Node> result3 = TreeUtil.buildTree(list);
        System.out.println("-----【buildTree】-----");
        System.out.println(JSON.toJSONString(result3));
    }


}