Java树迭代器（支持前序、后序和层级遍历）

1、设计背景

为简化实际业务场景中常见的菜单、组织架构等树型数据结构的处理，避免代码中出现相似代码和复杂代码，采用迭代器模式，设计一套树迭代器，支持业务场景中常用的前序遍历、后序遍历和层级遍历。

2、UML图

3、使用demo

1、SimpleTree

@Data
public class SimpleTree {

    /**
     * 节点名称
     */
    private String nodeName;

    /**
     * 子节点集合
     */
    private List<SimpleTree> children;
}

2、IteratorClient

public class IteratorClient {

    public static void main(String[] args) {
        Tree<SimpleTree> tree = new Tree<>(initTree(), SimpleTree::getChildren, SimpleTree::setChildren);
        TreeIterator<SimpleTree> iterator = tree.iterator(Tree.TreeIteratorEnum.LEVEL);
        while (iterator.hasNext()) {
            SimpleTree item = iterator.next();
            System.out.println(item.getNodeName());
            // 删除当前节点
            // iterator.remove();
            // 获取当前节点的父节点
            // SimpleTree currentParentNode = iterator.getCurrentParentNode();
            // 获取当前节点的子节点
            // List<SimpleTree> children = item.getChildren();
        }
    }

    /**
     * 组装树
     */
    private static List<SimpleTree> initTree() {
        List<SimpleTree> result = new ArrayList<>();
		// ... 填充树节点
        return result;
    }
}

4、树迭代器源码

1、TreeIterator

/**
 * 树迭代器-基础类
 *
 * @author mrdarren
 * @since 2021/7/8
 */
public abstract class TreeIterator<T> implements Iterator<T> {

    /**
     * 当前节点
     */
    T currentNode;

    /**
     * 当前父节点
     */
    T currentParentNode;

    /**
     * 子节点查询方法
     */
    Function<T, List<T>> getChildrenFun;

    TreeIterator(Function<T, List<T>> getChildrenFun) {
        this.getChildrenFun = getChildrenFun;
    }

    public T getCurrentParentNode() {
        return currentParentNode;
    }

    /**
     * 移除节点
     *
     * @author mrdarren
     * @since 2021/7/8
     */
    @Override
    public void remove() {
        // 判断当前父节点是否已取出
        if (this.currentParentNode != null) {
            // 获取子节点集合
            List<T> nodeList = getChildrenFun.apply(this.currentParentNode);
            // 移除当前子节点
            if (CollectionUtil.isNotEmpty(nodeList)) {
                nodeList.remove(this.currentNode);
            }
        }
    }
}

2、TreeLevelIterator

/**
 * 树-层级遍历迭代器
 *
 * @author mrdarren
 * @since 2021/7/8
 */
public class TreeLevelIterator<T> extends TreeIterator<T> {

    /**
     * 迭代队列
     */
    private Deque<T> queue = new LinkedList<>();

    /**
     * 父节点队列
     */
    private Deque<T> parentQueue = new LinkedList<>();

    public TreeLevelIterator(Function<T, List<T>> getChildrenFun, T root) {
        super(getChildrenFun);
        pushQueue(root);
    }

    /**
     * 压节点入队列
     *
     * @param root 父节点
     * @author mrdarren
     * @since 2021/7/8
     */
    private void pushQueue(T root) {
        if (root == null) {
            return;
        }
        List<T> nodeList = this.getChildrenFun.apply(root);
        if (CollectionUtil.isEmpty(nodeList)) {
            return;
        }
        nodeList.forEach(item -> {
            this.queue.add(item);
            this.parentQueue.add(root);
        });
    }

    /**
     * 是否有下一个节点
     *
     * @return {@link boolean}
     * @author mrdarren
     * @since 2021/7/8
     */
    @Override
    public boolean hasNext() {
        return !this.queue.isEmpty() && !this.parentQueue.isEmpty();
    }

    /**
     * 查询下一个节点
     *
     * @return {@link T}
     * @author mrdarren
     * @since 2021/7/8
     */
    @Override
    public T next() {
        if (!hasNext()) {
            throw new NoSuchElementException();
        }
        // pop双队列中第一个节点
        this.currentNode = this.queue.pop();
        this.currentParentNode = this.parentQueue.pop();

        // 子节点入队列
        pushQueue(this.currentNode);
        return this.currentNode;
    }
}

3、TreePostIterator

/**
 * 树-后序遍历迭代器
 * 遍历一次树，将前序遍历顺序存入队列用于外部迭代
 *
 * @author mrdarren
 * @since 2021/7/8
 */
public class TreePostIterator<T> extends TreeIterator<T> {

    /**
     * 迭代队列
     */
    private Deque<T> queue = new LinkedList<>();

    /**
     * 父节点队列
     */
    private Deque<T> parentQueue = new LinkedList<>();

    public TreePostIterator(Function<T, List<T>> getChildrenFun, T root) {
        super(getChildrenFun);
        initQueue(root);
    }

    /**
     * 是否有下一个节点
     *
     * @return {@link boolean}
     * @author mrdarren
     * @since 2021/7/8
     */
    @Override
    public boolean hasNext() {
        return !this.queue.isEmpty() && !this.parentQueue.isEmpty();
    }

    /**
     * 查询下一个节点
     *
     * @return {@link T}
     * @author mrdarren
     * @since 2021/7/8
     */
    @Override
    public T next() {
        if (!hasNext()) {
            throw new NoSuchElementException();
        }

        // pop双队列中的第一个节点
        this.currentNode = this.queue.pop();
        this.currentParentNode = this.parentQueue.pop();
        return this.currentNode;
    }

    /**
     * 初始化队列
     *
     * @param root 根节点
     * @author mrdarren
     * @since 2021/7/8
     */
    private void initQueue(T root) {
        if (root == null) {
            return;
        }

        List<T> nodeList = this.getChildrenFun.apply(root);
        if (CollectionUtil.isEmpty(nodeList)) {
            return;
        }

        // 创建临时栈，将二级节点入栈
        Stack<T> stack = new Stack<>();
        Stack<T> parentStack = new Stack<>();
        nodeList.forEach(item -> {
            stack.push(item);
            parentStack.push(root);
        });

        // 遍历树
        while (!stack.isEmpty()) {
            // pop子节点和对应的父节点
            T node = stack.pop();
            T parentNode = parentStack.pop();
            // 子节点和父节点入队列
            this.queue.push(node);
            this.parentQueue.push(parentNode);

            List<T> subNodeList = this.getChildrenFun.apply(node);
            if (CollectionUtil.isNotEmpty(subNodeList)) {
                subNodeList.forEach(item -> {
                    stack.push(item);
                    parentStack.push(node);
                });
            }
        }
    }
}

4、TreePreIterator

/**
 * 树-前序遍历迭代器
 *
 * @author mrdarren
 * @since 2021/7/8
 */
public class TreePreIterator<T> extends TreeIterator<T> {

    /**
     * 迭代栈
     */
    private Stack<T> stack = new Stack<>();

    /**
     * 父节点栈
     */
    private Stack<T> parentStack = new Stack<>();

    public TreePreIterator(Function<T, List<T>> getChildrenFun, T root) {
        super(getChildrenFun);
        pushChildrenToStack(root);
    }

    /**
     * 是否有下一个节点
     *
     * @return {@link boolean}
     * @author mrdarren
     * @since 2021/7/8
     */
    @Override
    public boolean hasNext() {
        return !this.stack.isEmpty() && !this.parentStack.isEmpty();
    }

    /**
     * 查询下一个节点
     *
     * @return {@link T}
     * @author mrdarren
     * @since 2021/7/8
     */
    @Override
    public T next() {
        if (!hasNext()) {
            throw new NoSuchElementException();
        }
        // pop双栈中最后一个节点
        this.currentNode = this.stack.pop();
        this.currentParentNode = this.parentStack.pop();

        // 子节点入栈
        pushChildrenToStack(this.currentNode);
        return this.currentNode;
    }

    /**
     * 入栈
     *
     * @param node 节点
     * @author mrdarren
     * @since 2021/7/8
     */
    private void pushChildrenToStack(T node) {
        if (node == null) {
            return;
        }
        // 查询子节点
        List<T> nodeList = this.getChildrenFun.apply(node);
        if (CollectionUtil.isEmpty(nodeList)) {
            return;
        }
        // 先压右子节点，再压左子节点
        for (int i = nodeList.size() - 1; i >= 0; i--) {
            // 子节点入栈
            this.stack.push(nodeList.get(i));
            // 子节点对应的父节点入栈
            this.parentStack.push(node);
        }
    }
}

5、Tree

/**
 * 树对象
 *
 * @author mrdarren
 * @since 2021/7/8
 */
public class Tree<T> {

    /**
     * 根节点集合
     */
    private T root;

    /**
     * 子节点查询方法
     */
    private Function<T, List<T>> getChildrenFun;

    /**
     * 构造方法，将树转为统一根对象
     *
     * @param rootList       根节点集合
     * @param getChildrenFun 获取子节点方法
     * @param setChildrenFun 设置子节点方法
     * @author mrdarren
     * @since 2021/7/8
     */
    @SuppressWarnings("unchecked")
    public Tree(List<T> rootList, Function<T, List<T>> getChildrenFun, BiConsumer<T, List<T>> setChildrenFun) {
        try {
            this.getChildrenFun = getChildrenFun;
            if (CollectionUtil.isNotEmpty(rootList)) {
                T firstRoot = rootList.get(0);
                T newRoot = (T) firstRoot.getClass().newInstance();
                setChildrenFun.accept(newRoot, rootList);
                this.root = newRoot;
            } else {
                this.root = null;
            }
        } catch (Exception e) {
            e.printStackTrace();
            throw new ClassCastException();
        }
    }

    /**
     * 根据类型返回对应的迭代器
     *
     * @param treeIteratorEnum 迭代器类型
     * @return {@link TreeIterator}
     * @author mrdarren
     * @since 2021/7/8
     */
    public TreeIterator<T> iterator(TreeIteratorEnum treeIteratorEnum) {
        TreeIterator<T> iterator;
        switch (treeIteratorEnum) {
            case PRE:
                iterator = new TreePreIterator<>(this.getChildrenFun, this.root);
                break;
            case LEVEL:
                iterator = new TreeLevelIterator<>(this.getChildrenFun, this.root);
                break;
            case POST:
                iterator = new TreePostIterator<>(this.getChildrenFun, this.root);
                break;
            default:
                iterator = new TreePreIterator<>(this.getChildrenFun, this.root);
                break;
        }
        return iterator;
    }

    /**
     * 迭代器类型枚举
     *
     * @author mrdarren
     * @since 2021/7/8
     */
    public enum TreeIteratorEnum {
        /**
         * 前序遍历
         */
        PRE,
        /**
         * 后序遍历
         */
        POST,
        /**
         * 层级遍历
         */
        LEVEL
    }
}