LeetCode 341 - Flatten Nested List Iterator

作者 QIFAN 日期 2017-02-17
LeetCode 341 - Flatten Nested List Iterator

原题链接: 341. Flatten Nested List Iterator


题干

给定一个有很多层嵌套的数字列表,设计一个 Iterator 用来展开所有元素。

举例
输入: [[1,1],2,[1,1]]
输出:[1,1,2,1,1]

思路

当栈不为空,pop 元素,如果是 Integer ,插入 res ,如果是 List ,则倒序将元素 push 进 stack 进行下一轮循环。

代码:

public class NestedIterator implements Iterator<Integer> {
Stack<NestedInteger> listStack;
public NestedIterator(List<NestedInteger> nestedList) {
listStack = new Stack<>();
int size = nestedList.size();
for (int i = size - 1; i >= 0; i--) {
listStack.push(nestedList.get(i));
}
}
@Override
public Integer next() {
return listStack.pop().getInteger();
}
@Override
public boolean hasNext() {
while (!listStack.isEmpty()) {
NestedInteger curr = listStack.peek();
if (curr.isInteger()) {
return true;
}
listStack.pop();
int size = curr.getList().size();
for (int i = size - 1; i >= 0; i--) {
listStack.push(curr.getList().get(i));
}
}
return false;
}
}