package cn.temptation.util;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * 树工具类
 */
public class TreeUtil {
    /**
     * Tree装配方法
     *
     * @param tempTreeNodes
     * @param treeNodes
     * @return
     */
    public static List<TreeNode> Assemble(List<TreeNode> tempTreeNodes, List<TreeNode> treeNodes) {
        if (tempTreeNodes != null) {
            Map<Integer, TreeNode> map = new LinkedHashMap<>();
            for (TreeNode tn : tempTreeNodes) {
                map.put(tn.getId(), tn);
            }

            TreeNode treeNode;
            TreeNode pTreeNode;
            for (Integer id : map.keySet()) {
                treeNode = map.get(id);
                if (treeNode.getParentId() == 0) {
                    treeNodes.add(treeNode);
                } else {
                    pTreeNode = map.get(treeNode.getParentId());
                    List<TreeNode> children = pTreeNode.getChildren();
                    if (children != null) {
                        children.add(treeNode);
                    } else {
                        children = new ArrayList();
                        children.add(treeNode);
                        pTreeNode.setChildren(children);
                    }
                }
            }
        }

        return treeNodes;
    }
}