Flatten Nested List Iterator
Stack of reverse-iterators; unwind lists lazily so each next() is amortized O(1).
4 min read
stack iterator design
Problem#
Implement an iterator over a nested integer list. The iterator should support next() and hasNext() that flatten the structure.
Examples#
[[1,1],2,[1,1]]→[1,1,2,1,1][1,[4,[6]]]→[1,4,6]
Constraints#
1 <= nestedList.length <= 500. Values fit in 32-bit ints.
Approach#
Maintain a stack of pairs (iterator, end). hasNext() first peels any top frame whose iterator hit end, then if the current element is a list, push its iterators and recurse. Once the top frame’s current element is an integer, return true.
Solution#
class NestedIterator { stack<pair<vector<NestedInteger>::iterator, vector<NestedInteger>::iterator>> st;public: NestedIterator(vector<NestedInteger>& nestedList) { st.push({nestedList.begin(), nestedList.end()}); }
int next() { hasNext(); return (st.top().first++)->getInteger(); }
bool hasNext() { while (!st.empty()) { auto& [it, end] = st.top(); if (it == end) { st.pop(); continue; } if (it->isInteger()) return true; auto& sub = it->getList(); ++it; st.push({sub.begin(), sub.end()}); } return false; }};// NestedInteger interface is provided by the LeetCode runtime:// IsInteger() bool; GetInteger() int; GetList() []*NestedInteger
type frame struct { list []*NestedInteger idx int}
type NestedIterator struct { stack []*frame}
func Constructor(nestedList []*NestedInteger) *NestedIterator { return &NestedIterator{stack: []*frame{{list: nestedList, idx: 0}}}}
func (it *NestedIterator) HasNext() bool { for len(it.stack) > 0 { top := it.stack[len(it.stack)-1] if top.idx == len(top.list) { it.stack = it.stack[:len(it.stack)-1] continue } cur := top.list[top.idx] if cur.IsInteger() { return true } top.idx++ it.stack = append(it.stack, &frame{list: cur.GetList(), idx: 0}) } return false}
func (it *NestedIterator) Next() int { it.HasNext() top := it.stack[len(it.stack)-1] v := top.list[top.idx].GetInteger() top.idx++ return v}from typing import List
# NestedInteger interface is provided by the LeetCode runtime:# isInteger() -> bool; getInteger() -> int; getList() -> List[NestedInteger]
class NestedIterator: def __init__(self, nestedList: List['NestedInteger']): self.stack: List[list] = [[nestedList, 0]]
def hasNext(self) -> bool: while self.stack: top = self.stack[-1] lst, idx = top[0], top[1] if idx == len(lst): self.stack.pop() continue cur = lst[idx] if cur.isInteger(): return True top[1] += 1 self.stack.append([cur.getList(), 0]) return False
def next(self) -> int: self.hasNext() top = self.stack[-1] v = top[0][top[1]].getInteger() top[1] += 1 return v// NestedInteger interface is provided by the LeetCode runtime:// isInteger() -> bool; getInteger() -> int; getList() -> NestedInteger[]
var NestedIterator = function(nestedList) { this.stack = [[nestedList, 0]];};
NestedIterator.prototype.hasNext = function() { while (this.stack.length) { const top = this.stack[this.stack.length - 1]; const [lst, idx] = top; if (idx === lst.length) { this.stack.pop(); continue; } const cur = lst[idx]; if (cur.isInteger()) return true; top[1]++; this.stack.push([cur.getList(), 0]); } return false;};
NestedIterator.prototype.next = function() { this.hasNext(); const top = this.stack[this.stack.length - 1]; const v = top[0][top[1]].getInteger(); top[1]++; return v;};// NestedInteger interface is provided by the LeetCode runtime:// boolean isInteger(); Integer getInteger(); List<NestedInteger> getList();
public class NestedIterator implements Iterator<Integer> { private Deque<int[]> idxStack; private Deque<List<NestedInteger>> listStack;
public NestedIterator(List<NestedInteger> nestedList) { idxStack = new ArrayDeque<>(); listStack = new ArrayDeque<>(); listStack.push(nestedList); idxStack.push(new int[]{0}); }
@Override public boolean hasNext() { while (!listStack.isEmpty()) { List<NestedInteger> lst = listStack.peek(); int[] idx = idxStack.peek(); if (idx[0] == lst.size()) { listStack.pop(); idxStack.pop(); continue; } NestedInteger cur = lst.get(idx[0]); if (cur.isInteger()) return true; idx[0]++; listStack.push(cur.getList()); idxStack.push(new int[]{0}); } return false; }
@Override public Integer next() { hasNext(); List<NestedInteger> lst = listStack.peek(); int[] idx = idxStack.peek(); int v = lst.get(idx[0]).getInteger(); idx[0]++; return v; }}// NestedInteger interface is provided by the LeetCode runtime:// isInteger(): boolean; getInteger(): number; getList(): NestedInteger[]
class NestedIterator { private stack: Array<[NestedInteger[], number]> = [];
constructor(nestedList: NestedInteger[]) { this.stack.push([nestedList, 0]); }
hasNext(): boolean { while (this.stack.length) { const top = this.stack[this.stack.length - 1]; const [lst, idx] = top; if (idx === lst.length) { this.stack.pop(); continue; } const cur = lst[idx]; if (cur.isInteger()) return true; top[1]++; this.stack.push([cur.getList(), 0]); } return false; }
next(): number { this.hasNext(); const top = this.stack[this.stack.length - 1]; const v = top[0][top[1]].getInteger(); top[1]++; return v; }}Editorial#
The iterator pair (current, end) lets us pause and resume each list. Pushing a sub-list’s iterators only when we encounter it (in hasNext) keeps the work amortized: each element is touched a number of times proportional to its nesting depth, which sums to the total node count.
Complexity#
- Time: O(N) total over all calls; each individual call is amortized O(1).
- Space: O(D) where D is maximum nesting depth.
Concept revision#
Related#