Medium | LeetCode 341. 扁平化嵌套列表迭代器 | 递归 | 栈

341. 扁平化嵌套列表迭代器

给你一个嵌套的整型列表。请你设计一个迭代器，使其能够遍历这个整型列表中的所有整数。

列表中的每一项或者为一个整数，或者是另一个列表。其中列表的元素也可能是整数或是其他列表。

示例 1:

输入: [[1,1],2,[1,1]]
输出: [1,1,2,1,1]
解释: 通过重复调用 next 直到 hasNext 返回 false，next 返回的元素的顺序应该是: [1,1,2,1,1]。

示例 2:

输入: [1,[4,[6]]]
输出: [1,4,6]
解释: 通过重复调用 next 直到 hasNext 返回 false，next 返回的元素的顺序应该是: [1,4,6]。

解题思路

方法一： 递归

类似一个树, 可以采用深度优先的方法将一棵树的元素遍历, 然后将其放到List当中。然后使用List的迭代器做判断。

public static class NestedIterator implements Iterator<Integer> {
    private List<Integer> vals;
    private Iterator<Integer> cur;

    public NestedIterator(List<NestedInteger> nestedList) {
        vals = new ArrayList<Integer>();
        // 在构造对象的过程, 首先将嵌套的列表做处理, 直接将他拍平, 放到java.util.List当中
        dfs(nestedList);
        cur = vals.iterator();
    }

    @Override
    public Integer next() {
        return cur.next();
    }

    @Override
    public boolean hasNext() {
        return cur.hasNext();
    }

    private void dfs(List<NestedInteger> nestedList) {
        for (NestedInteger nest : nestedList) {
            if (nest.isInteger()) {
                vals.add(nest.getInteger());
            } else {
                dfs(nest.getList());
            }
        }
    }
}

方法二： 栈

如上能够使用递归, 所以用栈也能实现

class NestedIterator2 implements Iterator<Integer> {
    // 存储列表的当前遍历位置
    private Deque<Iterator<NestedInteger>> stack;

    public NestedIterator2(List<NestedInteger> nestedList) {
        stack = new LinkedList<Iterator<NestedInteger>>();
        stack.push(nestedList.iterator());
    }

    @Override
    public Integer next() {
        // 由于保证调用 next 之前会调用 hasNext，直接返回栈顶列表的当前元素
        // 这是由题目保证的
        return stack.peek().next().getInteger();
    }

    @Override
    public boolean hasNext() {
        while (!stack.isEmpty()) {
            // 先拿出栈顶的迭代器
            Iterator<NestedInteger> it = stack.peek();
            // 判断当前的迭代器是否是遍历完了
            if (!it.hasNext()) {
                // 如果遍历完了, 将当前的迭代器出栈,
                // 然后将在下一轮循环中拿到下一个栈顶的迭代器
                stack.pop();
                continue;
            }
            // 程序走到这里说明栈顶的迭代器没有遍历完所有的元素, 则直接取出下一个元素
            // 若取出的元素是整数，则通过创建一个额外的列表将其重新放入栈中
            NestedInteger nest = it.next();
            if (nest.isInteger()) {
                // 由于要把迭代器的数字next()之后才能判断是否是数字还是列表
                // 所以需要在判断完成之后, 就将其重新加入到栈当中
                List<NestedInteger> list = new ArrayList<NestedInteger>();
                list.add(nest);
                stack.push(list.iterator());
                return true;
            }
            stack.push(nest.getList().iterator());
        }
        return false;
    }
}