实现一个简单的TreeUtils

码畜c

已于 2022-12-28 13:59:20 修改

阅读量114

点赞数

分类专栏： Java 文章标签： java

于 2022-12-11 15:46:39 首次发布

本文链接：https://blog.csdn.net/qq_38074398/article/details/128232572

版权

Java 专栏收录该内容

21 篇文章 0 订阅

订阅专栏

在b站上刷到了一个实现类似功能的视频，感觉不是特别灵活。故也自实现了一版。

代码

import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.Deque;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.function.Predicate;

/**
 * Created by Kim QQ.Cong on 2022-04-24 / 21:18
 *
 * @author: CongQingquan
 * @Description: 通用树工具类
 */
public class TreeUtils {

    private TreeUtils() {
    }

    /**
     * 树化方法的参数封装接口
     *
     * @param <T>
     * @param <ID>
     */
    public interface ToTreeParam<T, ID> {
    
        /**
         * 获取id提取器 (id & pid同类型)
         */
        Function<T, ID> getIdExtractor();

        /**
         * 获取pid提取器 (id & pid同类型)
         */
        Function<T, ID> getPidExtractor();
        
        /**
         * 获取节点排序比较器
         */
        Comparator<? super T> getComparator();
        
        /**
         * 获取设置子节点的BiConsumer
         */
        BiConsumer<T, Collection<T>> getSettingChildrenBiConsumer();
        
        /**
         * 获取默认的根元素断言 (test: pid == null)
         */
        default Predicate<T> getRootPredicate() {
            return (n) -> this.getPidExtractor().apply(n) == null;
        }
    }

    /**
     * 树化
     *
     * @param param 参数封装类
     * @param <T>
     * @param <ID>
     * @return
     */
    public static <T, ID> Collection<T> toTree(Collection<T> nodes, ToTreeParam<T, ID> param) {
        if (param == null) {
            throw new IllegalArgumentException("To tree param cannot be null");
        }
        return toTree(nodes, param.getIdExtractor(), param.getPidExtractor(), param.getRootPredicate(),
                param.getComparator(), param.getSettingChildrenBiConsumer());
    }

    /**
     * 树化
     * 1. id data type === parent id data type
     * 2. 默认的，root node: pidExtractor.apply(node) == null & child node: pidExtractor.apply(node) != null
     *
     * @param nodes           节点集合
     * @param idExtractor     id提取器 (id & pid同类型)
     * @param pidExtractor    pid提取器 (id & pid同类型)
     * @param <T>             节点类型
     * @param <ID>            id数据类型 (id & pid同类型)
     * @param isRoot          根节点断言
     * @param comparator      节点排序比较器
     * @param settingChildren 设置子节点
     *                        1. 为什么用 set 的方式而不是 get 的方式？
     *                           考虑 get 的 children collection 可能为 null 的情况，初始化后无法进行 children 数据的写回，希望调用者提前写初始化操作
     *                        2. 调用者在初始化 children 容器时，根据容器底层使用的数据结构，可能会影响到 children 节点根据 comparator 的排序顺序
     *
     * @return
     */
    public static <T, ID> Collection<T> toTree(final Collection<T> nodes,
                                               final Function<T, ID> idExtractor,
                                               final Function<T, ID> pidExtractor,
                                               final Predicate<T> isRoot,
                                               final Comparator<? super T> comparator,
                                               final BiConsumer<T, Collection<T>> settingChildren) {
        if (nodes == null) {
            return new ArrayList<>();
        }
        if (nodes.size() == 0) {
            return nodes;
        }
        if (idExtractor == null || pidExtractor == null) {
            throw new IllegalArgumentException("Extractor cannot be null");
        }
        if (isRoot == null) {
            throw new IllegalArgumentException("Root node predicate cannot be null");
        }
        final Set<T> roots = new TreeSet<>(comparator);
        final Map<ID, Collection<T>> childrenMapping = new HashMap<>(16);
        nodes.forEach(node -> {
            if (isRoot.test(node)) {
                roots.add(node);
                return;
            }
            childrenMapping.computeIfAbsent(pidExtractor.apply(node), (id) -> new TreeSet<>(comparator)).add(node);
        });
    
        final Deque<T> stack = new LinkedList<>();
        roots.forEach(root -> {
            stack.add(root);
            while (stack.size() > 0) {
                T currentNode = stack.pop();
                ID currentNodeId = idExtractor.apply(currentNode);
                Collection<T> cc = childrenMapping.get(currentNodeId);
                if (CollectionUtils.isEmpty(cc)) {
                    continue;
                }
                for (T cn : cc) {
                    stack.push(cn);
                }
                settingChildren.accept(currentNode, cc);
            }
        });
        return roots;
    }
}

代码设计

利用Stack而非递归的形式实现的子节点搜索，并且节点类不需要继承任何类似Node这样的内部提供的Base类。一般来说可能需要通过定义一个抽象的Node类，里面定义好最基本的字段，如：id、pid、children。并提供对应的set/get方法。但是这样需要强制节点类继承Node类，如果实体类已有继承类呢？如果实体类已经有了其中的某些字段呢？约束过多，体验不好。所以我想通过lambda传参的形式，在配合泛型，来完成这些操作。

比如：

id & pid 的提取器，获取实体类的id & pid字段值。这里的id&pid是抽象的，具体是实体类的哪个字段提供，由调用者决定。
isRoot 的断言，这是刚才没有提到的。如何判断节点为一级根节点，不能仅仅通过常规的pid==null的形式来判定，这无疑又进行了一次强制约束。比如，我就是想要将二级节点作为根节点，且节点的pid就是不为空。所以，还是交给调用者来决定。
settingChildren，为什么设定children也被提取出来？有两层含义：
1. 调用者来确定children对应实体类的哪个字段
2. 不确定获取的实体类的children集合字段是否为null，如果为null，那么无法进行children数据的写回

参数简化

动态的代价基本都是实现代码的增加或调用代码增加。如：

Collection<Node> nodes = TreeUtils.toTree(roots, Node::getId, Node::getParentId,
        (node) -> node.getParentId() == null,
        Comparator.comparing(Node::getId).reversed(),
        (node, cns) -> {
            List<Node> cl = node.getChildren();
            if (cl == null) {
                node.setChildren(Lists.newArrayList(cns));
                return;
            }
            cl.addAll(cns);
        });

那么为了进行简化，内部提供了参数封装接口：ToTreeParam。简化后的代码如下：

class DefaultToTreeParam implements ToTreeParam<Node, String> {

    @Override
    public Function<Node, String> getIdExtractor() {
        return Node::getId;
    }

    @Override
    public Function<Node, String> getPidExtractor() {
        return Node::getParentId;
    }

    @Override
    public Predicate<Node> getRootPredicate() {
        return (n) -> n.getId().length() == 2;
    }

    @Override
    public Comparator<? super Node> getComparator() {
        return Comparator.comparing(Node::getId).reversed();
    }

    @Override
    public BiConsumer<Node, Collection<Node>> getSettingChildrenBiConsumer() {
        return (node, cns) -> {
            List<Node> cl = node.getChildren();
            if (cl == null) {
                node.setChildren(Lists.newArrayList(cns));
                return;
            }
            cl.addAll(cns);
        };
    }
}

Collection<Node> nodes2 = TreeUtils.toTree(data, new DefaultToTreeParam());