工作总结—集合转树与树的搜索

您需要了解

1. Java 泛型（上下界，泛型方法）
2. stream 操作
3. 递归
4. 函数式接口

运行环境

java version "1.8.0_202"

Show me the code

定义树结构

package top.mindse.cloud;

import java.io.Serializable;
import java.util.*;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;

public class TreeUtils {
    /**
     * 树结构
     */
    public static class TreeNode implements Serializable {
        private String id;
        private List<? extends TreeNode> child;
        private String parentId;

        public TreeNode() {
        }

        public TreeNode(String id, String parentId) {
            this.id = id;
            this.parentId = parentId;
        }

        public String getId() {
            return id;
        }

        public void setId(String id) {
            this.id = id;
        }

        public List<? extends TreeNode> getChild() {
            return child;
        }

        public void setChild(List<? extends TreeNode> child) {
            this.child = child;
        }

        public String getParentId() {
            return parentId;
        }

        public void setParentId(String parentId) {
            this.parentId = parentId;
        }
        
        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            TreeNode treeNode = (TreeNode) o;
            return Objects.equals(id, treeNode.id);
        }

        @Override
        public int hashCode() {
            return Objects.hash(id);
        }
    }
}

树化

  /**
     * 树化
     *
     * @param source
     * @param <T>
     * @return
     */
    public static <T extends TreeNode> List<T> tree(List<T> source) {
        if (Objects.isNull(source) || source.isEmpty()) {
            return new ArrayList<>();
        }
        Map<String, T> map = source.stream().collect(Collectors.toMap(TreeNode::getId, Function.identity()));
        Set<T> root = new HashSet<>();
        source.forEach(d -> {
            String parentId = d.getParentId();
            if (map.containsKey(parentId)) {
                T parent = map.get(parentId);
                List<T> child = (List<T>) parent.getChild();
                if (child == null) {
                    child = new ArrayList<>();
                }
                child.add(d);
                parent.setChild(child);
                root.remove(d);
            } else {
                root.add(d);
            }
        });
        return new ArrayList<>(root);
    }

树的搜索

 /**
     * 搜索树
     *
     * @param source
     * @param predicate
     * @param <T>
     * @return
     */
    public static <T extends TreeNode> List<T> search(List<T> source, Predicate<T> predicate) {
        List<T> target = new ArrayList<>();
        source.forEach(o -> {
            List<T> child = (List<T>) o.getChild();
            if (Objects.nonNull(child) && !child.isEmpty()) {
                List<T> searched = search(child, predicate);
                if (!searched.isEmpty()) {
                    o.setChild(searched);
                    target.add(o);
                } else {
                    if (predicate.test(o)) {
                        target.add(o);
                    }
                }
            } else {
                if (predicate.test(o)) {
                    target.add(o);
                }
            }
        });
        return target;
    }

使用

package top.mindse.cloud;

import java.io.Serializable;
import java.util.*;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;

public class Solution {

    /**
     * 调用示例
     * @param args
     */
    public static void main(String[] args) {

        class BTreeNode extends TreeNode {

            public BTreeNode() {
            }

            public BTreeNode(String id, String parentId, String name) {
                super(id, parentId);
                this.name = name;
            }

            private String name;

            public String getName() {
                return name;
            }

            public void setName(String name) {
                this.name = name;
            }
        }

        String id = "12";
        List<BTreeNode> source = new ArrayList<>();
        BTreeNode l11 = new BTreeNode("1", "0", "第一级");
        source.add(l11);
        BTreeNode l111 = new BTreeNode("11", "1", "第一1级");
        source.add(l111);
        BTreeNode l112 = new BTreeNode("12", "1", "第一2级");
        source.add(l112);
        BTreeNode l12 = new BTreeNode("2", "0", "第二级");
        source.add(l12);
        List<BTreeNode> dataTree = tree(source);
        List<BTreeNode> rsp = search(dataTree, bTreeNode -> {
            if (Objects.equals(id, bTreeNode.getId())) {
                return true;
            }
            return false;
        });
        //自定义处理
    }
}