原题链接: 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; } }
|