leetCode-341:Flatten Nested List Iterator

问题描述

给定一个包含整数或整数列表的列表，要求使用迭代器将列表中的数字进行平铺。题目链接：**点我**

样例输入输出

输入：nestedList = [[1,1],2,[1,1]]

输出：[1,1,2,1,1]

输入：nestedList = [1,[4,[6]]]

输出：[1,4,6]

问题解法

此题由于有嵌套列表的存在，所以考虑用栈来保存嵌套的列表，由于遍历列表中需要用到下标，所以将 NestedInteger 和遍历的下标封装成对象。由于 [[]] 的存在，所以部分取数放数的逻辑放在 hasNext 下。代码如下

/**
 * // This is the interface that allows for creating nested lists.
 * // You should not implement it, or speculate about its implementation
 * public interface NestedInteger {
 *
 *     // @return true if this NestedInteger holds a single integer, rather than a nested list.
 *     public boolean isInteger();
 *
 *     // @return the single integer that this NestedInteger holds, if it holds a single integer
 *     // Return null if this NestedInteger holds a nested list
 *     public Integer getInteger();
 *
 *     // @return the nested list that this NestedInteger holds, if it holds a nested list
 *     // Return empty list if this NestedInteger holds a single integer
 *     public List<NestedInteger> getList();
 * }
 */
public class NestedIterator implements Iterator<Integer> {

    class Node {
        NestedInteger nestedInteger;
        int index;

        public Node(NestedInteger nestedInteger, int index) {
            this.nestedInteger = nestedInteger;
            this.index = index;
        }
    }

    private List<NestedInteger> nestedList;

    private Stack<Node> stack;

    public NestedIterator(List<NestedInteger> nestedList) {
        this.nestedList = nestedList;
        stack = new Stack<>();
        stack.push(new Node(null, 0));
    }

    @Override
    public Integer next() {
        Integer result = stack.pop().nestedInteger.getInteger();
        stack.peek().index++;
        return result;
    }

    @Override
    public boolean hasNext() {
        while (!stack.isEmpty()) {
            Node current = stack.peek();
            if (current.nestedInteger == null) {
                if (current.index == nestedList.size()) {
                    return false;
                }

                pushNode(nestedList.get(current.index));
                continue;
            }

            if (current.index == -1) {
                return true;
            }

            List<NestedInteger> list = current.nestedInteger.getList();
            if (current.index < list.size()) {
                NestedInteger temp = list.get(current.index);
                pushNode(temp);
            } else {
                stack.pop();
                stack.peek().index++;
            }
        }

        return false;
    }

    private void pushNode(NestedInteger nestedInteger) {
        int index = 0;
        if (nestedInteger.isInteger()) {
            index = -1;
        }

        stack.push(new Node(nestedInteger, index));
    }
}

/**
 * Your NestedIterator object will be instantiated and called as such:
 * NestedIterator i = new NestedIterator(nestedList);
 * while (i.hasNext()) v[f()] = i.next();
 */