TreeUtil

点击查看代码

 
import org.apache.commons.collections4.CollectionUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;
import java.util.function.Function;

/**
 * List -> TreeList
 */
public class TreeUtil<T, K> {

    /**
     * 普通的非树形 list
     */
    private List<T> list;
    /**
     * 树形 list
     */
    private List<T> treeList;
    /**
     * 节点的 id 的 getter 方法
     */
    private Function<T, K> idGetterFunction;
    /**
     * 节点的父 id 的 getter 方法
     */
    private Function<T, K> parentIdGetterFunction;
    /**
     * 节点的子集合的 setter 方法
     */
    private BiConsumer<T, List<T>> childListSetterConsumer;
    /**
     * 节点的子集合的 getter 方法
     */
    private Function<T, List<T>> childListGetterFunction;
    /**
     * 节点的层数的 setter 方法(根节点层数为 1,次根节点的层数为 2,以此类推)
     */
    private FloorConsumer<T, Integer> floorSetterConsumer;

    /**
     * 构造方法(不设置层数)
     */
    public TreeUtil(List<T> list,
                    Function<T, K> idGetterFunction,
                    Function<T, K> parentIdGetterFunction,
                    BiConsumer<T, List<T>> childListSetterConsumer) {
        this.list = list;
        this.idGetterFunction = idGetterFunction;
        this.parentIdGetterFunction = parentIdGetterFunction;
        this.childListSetterConsumer = childListSetterConsumer;
    }

    /**
     * 构造方法(设置层数)
     */
    public TreeUtil(List<T> list,
                    Function<T, K> idGetterFunction,
                    Function<T, K> parentIdGetterFunction,
                    BiConsumer<T, List<T>> childListSetterConsumer,
                    FloorConsumer<T, Integer> floorSetterConsumer) {
        this.list = list;
        this.idGetterFunction = idGetterFunction;
        this.parentIdGetterFunction = parentIdGetterFunction;
        this.childListSetterConsumer = childListSetterConsumer;
        this.floorSetterConsumer = floorSetterConsumer;
    }

    /**
     * 构造方法
     */
    public TreeUtil(List<T> treeList, Function<T, List<T>> childListGetterFunction) {
        this.treeList = treeList;
        this.childListGetterFunction = childListGetterFunction;
    }

    /**
     * 普通 list -> 树形 list
     */
    public List<T> fromListToTreeList() {
        if (list.isEmpty()) {
            return new ArrayList<>();
        }
        // 根节点集合
        List<T> rootList = new ArrayList<>();
        AtomicInteger atomicInteger = new AtomicInteger(1);
        for (T node : list) {
            // 若每个节点的 id 都不是这个节点的父id,则这个节点是根节点
            if (list.stream().noneMatch(i -> idGetterFunction.apply(i).equals(parentIdGetterFunction.apply(node)))) {
                rootList.add(buildTreeOfTheNode(node, 1, atomicInteger));
            }
        }
        return rootList;
    }

    /**
     * 树形 list -> 普通 list
     */
    public List<T> fromTreeListToList() {
        List<T> list = new ArrayList<>();
        if (treeList.isEmpty()) {
            return list;
        }
        // 添加根节点
        list.addAll(treeList);
        // 某个层数的所有节点
        List<T> listOfAFloor = new CopyOnWriteArrayList<>(treeList); // listOfAFloor: while 循环的入参
        while (true) {
            if (CollectionUtils.isEmpty(listOfAFloor)) {
                break;
            }
            // 某个层数的所有节点的子节点(不包括孙子节点)
            List<T> childListOfAFloor = new CopyOnWriteArrayList<>();
            for (T t : listOfAFloor) {
                List<T> childList = childListGetterFunction.apply(t);
                if (CollectionUtils.isNotEmpty(childList)) {
                    list.addAll(childList);
                    childListOfAFloor.addAll(childList);
                }
            }
            listOfAFloor = childListOfAFloor; // listOfAFloor: 赋值为 childListOfAFloor,作为下一个 while 循环的入参
        }
        return list;
    }

    /**
     * 通过递归,将某个节点和他的子,孙,曾孙 ... 节点构建为树形结构
     *
     * @param node  节点
     * @param floor 层数
     * @return 树形结构
     */
    private T buildTreeOfTheNode(T node, Integer floor, AtomicInteger atomicInteger) {
        int count = atomicInteger.getAndIncrement();
        if (count > 500) {
            throw new RuntimeException("递归次数过多");
        }
        List<T> childList = new ArrayList<>();
        if (floorSetterConsumer != null) {
            floorSetterConsumer.setNodeFloor(node, floor);
        }
        for (T t : list) {
            if (idGetterFunction.apply(node).equals(parentIdGetterFunction.apply(t))) {
                childList.add(buildTreeOfTheNode(t, (floor + 1), atomicInteger));
            }
        }
        childListSetterConsumer.accept(node, childList);
        return node;
    }

}

点击查看代码

@Data
class Person {
    /**
     * id
     */
    private Integer id;
    /**
     * 父id
     */
    private Integer pId;
    /**
     * 名称
     */
    private String name;
    /**
     * 子节点集合
     */
    private List<Person> sons;
    /**
     * 层数
     */
    private Integer floor;
}

点击查看代码

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

public class MainTest {

    public static void main(String[] args) {
        List<Person> list = new ArrayList<>();

        Person grandFather = new Person();
        grandFather.setId(1);
        grandFather.setPId(0);
        grandFather.setName("爷爷");
        list.add(grandFather);

        Person father = new Person();
        father.setId(2);
        father.setPId(1);
        father.setName("爸爸");
        list.add(father);

        Person uncle_1 = new Person();
        uncle_1.setId(3);
        uncle_1.setPId(1);
        uncle_1.setName("大叔");
        list.add(uncle_1);

        Person uncle_2 = new Person();
        uncle_2.setId(4);
        uncle_2.setPId(1);
        uncle_2.setName("二叔");
        list.add(uncle_2);

        Person grandson = new Person();
        grandson.setId(5);
        grandson.setPId(2);
        grandson.setName("小明");
        list.add(grandson);

        //TreeUtil<Node, Integer> treeUtil = new TreeUtil<>(list, Node::getId, Node::getPId, Node::setSons);
        TreeUtil<Person, Integer> treeUtil = new TreeUtil<>(list, Person::getId, Person::getPId, Person::setSons, Person::setFloor);
        List<Person> treeList = treeUtil.fromListToTreeList();
        System.out.println(treeList);

        TreeUtil<Person, Integer> unBuildTreeUtil = new TreeUtil<>(treeList, Person::getSons);
        List<Person> personList = unBuildTreeUtil.fromTreeListToList();
        System.out.println(personList);

    }

}