Flatten Nested List Iterator — Stack-Based Lazy Expansion
Advertisement
Problem 270 · Flatten Nested List Iterator
Difficulty: Medium · Pattern: Stack Design
Implement hasNext() and next() for a nested list iterator.
Solutions
# Python — stack approach
class NestedIterator:
def __init__(self, nestedList):
self.stack = list(reversed(nestedList))
def next(self):
return self.stack.pop().getInteger()
def hasNext(self):
while self.stack:
top = self.stack[-1]
if top.isInteger(): return True
self.stack.pop()
self.stack.extend(reversed(top.getList()))
return False
// Java
class NestedIterator implements Iterator<Integer> {
Deque<NestedInteger> stack = new ArrayDeque<>();
public NestedIterator(List<NestedInteger> nestedList) {
for (int i=nestedList.size()-1;i>=0;i--) stack.push(nestedList.get(i));
}
public Integer next() { return stack.pop().getInteger(); }
public boolean hasNext() {
while (!stack.isEmpty()) {
NestedInteger top=stack.peek();
if (top.isInteger()) return true;
stack.pop();
List<NestedInteger> list=top.getList();
for (int i=list.size()-1;i>=0;i--) stack.push(list.get(i));
}
return false;
}
}
Complexity
- Time: O(1) amortized per call
- Space: O(depth * width)
Advertisement