[Medium] 341. Flatten Nested List Iterator
You are given a nested list of integers nestedList. Each element is either an integer or a list whose elements may also be integers or other lists. Implement an iterator to flatten it.
Implement NestedIterator:
NestedIterator(vector<NestedInteger> &nestedList)– initializes the iteratorint next()– returns the next integer in the flattened listbool hasNext()– returnstrueif there are still integers to iterate
Examples
Example 1:
Input: nestedList = [[1,1],2,[1,1]]
Output: [1,1,2,1,1]
Explanation: Flattening gives [1,1,2,1,1].
Example 2:
Input: nestedList = [1,[4,[6]]]
Output: [1,4,6]
Constraints
1 <= nestedList.length <= 500- Values are in the range
[-10^6, 10^6] - At most
10^5calls tonextandhasNext
Thinking Process
The Problem with Pre-flattening
We could recursively flatten the entire list in the constructor and iterate over the result. That works, but uses $O(n)$ extra space upfront and doesn’t take advantage of lazy evaluation – we might not need all elements.
Stack-Based Lazy Approach
Use a stack to simulate the recursive flattening on demand:
- Constructor: push elements onto the stack in reverse order (so the first element is on top)
hasNext: peek at the top – if it’s a list, expand it (pop, push children in reverse). Repeat until the top is an integer or the stack is emptynext: the top is guaranteed to be an integer (sincehasNextwas called first) – pop and return it
Why Push in Reverse?
A stack is LIFO. To process elements left-to-right, we push them right-to-left:
nestedList = [A, B, C]
Push: C, B, A → stack top = A ✓
Walk-through
Input: [[1,1], 2, [1,1]]
Constructor: push [1,1], 2, [1,1] in reverse
stack: [1,1] | 2 | [1,1] (top → [1,1])
hasNext(): top = [1,1] → list, expand: push 1, 1
stack: [1,1] | 2 | 1 | 1 (top → 1, integer ✓)
next(): return 1
hasNext(): top = 1 (integer ✓)
next(): return 1
hasNext(): top = 2 (integer ✓)
next(): return 2
hasNext(): top = [1,1] → list, expand: push 1, 1
stack: 1 | 1 (top → 1 ✓)
next(): return 1
hasNext(): top = 1 (integer ✓)
next(): return 1
hasNext(): stack empty → false
Solution: Stack-Based Lazy Iterator – $O(1)$ amortized
class NestedIterator {
public:
NestedIterator(vector<NestedInteger> &nestedList) {
for (int i = nestedList.size() - 1; i >= 0; --i) {
stk.push(nestedList[i]);
}
}
int next() {
int val = stk.top().getInteger();
stk.pop();
return val;
}
bool hasNext() {
while (!stk.empty()) {
NestedInteger curr = stk.top();
if (curr.isInteger()) {
return true;
}
stk.pop();
auto& lst = curr.getList();
for (int i = lst.size() - 1; i >= 0; --i) {
stk.push(lst[i]);
}
}
return false;
}
private:
stack<NestedInteger> stk;
};
Time: $O(1)$ amortized per next/hasNext call – each element is pushed and popped exactly once across all calls
Space: $O(D)$ where $D$ = maximum nesting depth (stack holds at most one “path” through the nesting)
Key Details
Why does hasNext do the flattening, not next?
The iterator contract requires hasNext to correctly report whether more elements exist. If the top of the stack is a nested list (possibly empty), hasNext must drill down to find the next actual integer – or determine there are none.
What about empty nested lists like [[], [1]]?
The while loop in hasNext handles this: [] gets popped and expanded to nothing, then the loop continues to process [1].
Common Mistakes
- Pushing elements in forward order (reverses the iteration order)
- Flattening in
nextinstead ofhasNext(breaks the contract when checking for empty nested lists) - Not handling deeply nested empty lists like
[[[[]]]]
Key Takeaways
- “Flatten a recursive structure lazily” = stack-based iterator
- Push in reverse order to maintain left-to-right traversal with a LIFO stack
- Doing the expansion in
hasNextrather thannextcorrectly handles empty nested lists and satisfies the iterator contract
Related Problems
- 385. Mini Parser – building nested structures
- 339. Nested List Weight Sum – DFS on nested lists
- 173. Binary Search Tree Iterator – stack-based tree iterator
- 281. Zigzag Iterator – iterator design