jsl-questionnaire-base/src/main/java/com/qs/serve/common/util/TreeUtil.java

package com.qs.serve.common.util;

import com.qs.serve.common.model.dto.TreeNode;
import lombok.experimental.UtilityClass;

import java.util.ArrayList;
import java.util.List;

/**
 * @author YenHex
 * @since 2022/3/1
 */
@UtilityClass
public class TreeUtil {

    public static final String DEFAULT_PID_STRING = "0";

    /**
     * 两层循环实现建树
     * @param treeNodes 传入的树节点列表
     * @param rootPid
     * @return
     * sorted(Comparator.comparingInt(DemoEntity::getSort))
     */
    public <T extends TreeNode> List<T> build(List<T> treeNodes, String rootPid) {
        List<T> trees = new ArrayList<>();
        for (T treeNode : treeNodes) {
            if (rootPid.equals(treeNode.getParentId())) {
                trees.add(treeNode);
            }
            for (T it : treeNodes) {
                if (it.getParentId().equals(treeNode.getId())) {
                    treeNode.addChildren(it);
                }
            }
        }
        return trees;
    }

    /**
     * 使用递归方法建树
     *
     * @param treeNodes
     * @return
     */
    public <T extends TreeNode> List<T> buildByRecursive(List<T> treeNodes, Object root) {
        List<T> trees = new ArrayList<T>();
        for (T treeNode : treeNodes) {
            if (root.equals(treeNode.getParentId())) {
                trees.add(findChildren(treeNode, treeNodes,1));
            }
        }
        return trees;
    }

    /**
     * 递归查找子节点
     *
     * @param treeNodes
     * @return
     */
    public <T extends TreeNode> T findChildren(T treeNode, List<T> treeNodes,int level) {
        for (T it : treeNodes) {
            if (treeNode.getId().equals(it.getParentId())) {
                if (treeNode.getChildren() == null) {
                    treeNode.setChildren(new ArrayList<>());
                }
                treeNode.addChildren(findChildren(it, treeNodes,level+1));
            }
        }
        treeNode.setNodeLevel(level);
        return treeNode;
    }

}