LeetCode 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]。

思路一

直接DFS遍历所有元素并存储到数组中，然后遍历数组中元素

代码

public class NestedIterator implements Iterator<Integer> {

    private List<Integer> vals;
    private Iterator<Integer> it;

    public NestedIterator(List<NestedInteger> nestedList) {
        vals = new ArrayList<Integer>();
        dfs(nestedList);
        it = vals.iterator();
    }

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

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

    // 深度遍历所有的元素放入一个列表中
    private void dfs(List<NestedInteger> nestedList) {
        for(NestedInteger nest : nestedList) {
            if(nest.isInteger()) {
                vals.add(nest.getInteger());
            }
            else dfs(nest.getList());
        }
    }
}

思路二

将第一层所有元素逆序放入栈中，栈中元素可能只有一个元素，也可能有多个元素，每次需要下一个元素时再进行展开；

代码

public class NestedIterator implements Iterator<Integer> {

    //全局的栈
    private Deque<NestedInteger> stack = new ArrayDeque<>();

    public NestedIterator(List<NestedInteger> nestedList) {
        //把第一层元素逆序放入栈中
        for(int i = nestedList.size() - 1; i >= 0; i--) {
            NestedInteger item = nestedList.get(i);
            stack.addLast(item);
        }
    }

    @Override
    public Integer next() {
        return hasNext() ? stack.pollLast().getInteger() : -1;
    }

    @Override
    public boolean hasNext() {
        //如果栈为空，则返回
        if(stack.isEmpty()) return false;
        else {
            //取栈顶元素
            NestedInteger item = stack.peekLast();
            if(item.isInteger()) return true;
            else {
                //若不是数字，则弹出栈顶元素
                item = stack.pollLast();
                List<NestedInteger> list = item.getList();
                //将这个子元素展开并逆序放入栈中
                for(int i = list.size() - 1; i >= 0; i --) {
                    stack.addLast(list.get(i));
                }
                //继续递归直到找到一个元素
                return hasNext();
            }
        }
    }
}