Problem
Given a nested list of integers, implement an iterator to flatten it.
Each element is either an integer, or a list -- whose elements may also be integers or other lists.
Example
Given the list [[1,1],2,[1,1]]
, By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1,1,2,1,1]
.
Given the list [1,[4,[6]]]
, By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1,4,6]
.
Note
这个迭代器要实现三个功能:
初始化结构,找下一个元素,判断是否存在下一个元素。
首先,根据迭代器需要不断返回下一个元素,确定用堆栈来做。
堆栈初始化数据结构,要先从后向前向堆栈压入nestedList
中的元素。
在调用next()
之前,先要用hasNext()
判断下一个NestedInteger
是Integer
还是List<Integer>
,并进行flatten的操作:对List<Integer>
要展开并顺序压入stack
;对Integer
直接返回true
。这样一来,调用next()
的时候,返回的就一定是Integer
了。
Solution
import java.util.Iterator;
public class NestedIterator implements Iterator<Integer> {
Stack<NestedInteger> stack = new Stack();
public NestedIterator(List<NestedInteger> nestedList) {
for (int i = nestedList.size()-1; i >= 0; i--) {
stack.push(nestedList.get(i));
}
}
@Override
public Integer next() {
return stack.pop().getInteger();
}
@Override
public boolean hasNext() {
while (!stack.isEmpty()) {
NestedInteger cur = stack.peek();
if (cur.isInteger()) return true;
stack.pop();
for (int i = cur.getList().size()-1; i >= 0; i--) stack.push(cur.getList().get(i));
}
return false;
}
}
By using internalNext()
public class NestedIterator implements Iterator<Integer> {
private NestedInteger peek = null;
private Iterator<NestedInteger> iterator;
private Stack<Iterator<NestedInteger>> stack = new Stack();
public NestedIterator(List<NestedInteger> nestedList) {
iterator = nestedList.iterator();
peek = internalNext();
}
private NestedInteger internalNext() {
if (iterator.hasNext()) {
NestedInteger i = iterator.next();
if (i.isInteger()) return i;
else {
stack.add(iterator);
iterator = i.getList().iterator();
return internalNext();
}
} else if (stack.size() > 0) {
iterator = stack.pop();
return internalNext();
} else return null;
}
@Override
public Integer next() {
Integer next = peek.getInteger();
peek = internalNext();
return next;
}
@Override
public boolean hasNext() {
return peek != null;
}
}
A much much better method
public class NestedIterator implements Iterator<Integer> {
private List<Integer> flatList;
private int index;
public NestedIterator(List<NestedInteger> nestedList) {
flatList = new ArrayList<>();
flatten(nestedList);
}
public void flatten(List<NestedInteger> nestedList) {
for (NestedInteger i: nestedList) {
if (i.isInteger()) flatList.add(i.getInteger());
else flatten(i.getList());
}
}
@Override
public Integer next() {
return hasNext ? flatList.get(index++) : null;
}
@Override
public boolean hasNext() {
return index < flatList.size();
}
}