DSA Stacks

Flatten Nested List Iterator

Stack of reverse-iterators; unwind lists lazily so each next() is amortized O(1).

Medium
Time O(N) amortized total Space O(D)
LeetCode
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#

Flatten Nested List Iterator
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;
}
};

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#

Search ESC

Keyboard shortcuts

Shortcuts are disabled while typing in inputs.