Flatten Nested List Iterator

The Problem

LeetCode 341 asks you to implement an iterator over a nested list of integers. Each element in the list is either an integer or a list whose elements may also be integers or other lists. You need to implement two methods:

• next() -- returns the next integer in the flattened sequence.
• hasNext() -- returns True if there are still integers remaining.

You are given a NestedInteger interface with three methods: isInteger(), getInteger(), and getList(). Your job is to wrap the nested structure and expose it as a flat iterator.

The diagram above shows [[1,1],2,[1,1]]. The iterator should produce 1, 1, 2, 1, 1 in that order, as though you had recursively unwrapped every nested list.

The Stack Approach

The naive approach is to flatten everything upfront in __init__ and then just iterate over the resulting list. That works and passes most test cases, but it does all the work before the caller even asks for a single element. If the caller only needs the first few integers, you have done unnecessary work.

A stack-based lazy approach does the right amount of work on demand:

1. In __init__, push all elements of the top-level list onto the stack in reverse order, so the first element sits at the top.
2. In hasNext(), peek at the top of the stack. If it is an integer, return True. If it is a list, pop it and push its elements in reverse. Repeat until the top is an integer or the stack is empty.
3. In next(), call hasNext() to guarantee the top is an integer, then pop and return it.

The key is that hasNext() does the structural work. When you encounter a nested list on top, you explode it one layer at a time. You never recurse into the whole tree at once.

The Solution

class NestedIterator:
    def __init__(self, nestedList):
        self.stack = list(reversed(nestedList))

    def next(self):
        self.hasNext()
        return self.stack.pop().getInteger()

    def hasNext(self):
        while self.stack:
            top = self.stack[-1]
            if top.isInteger():
                return True
            self.stack.pop()
            for item in reversed(top.getList()):
                self.stack.append(item)
        return False

A few things to notice:

• list(reversed(nestedList)) creates the initial stack with the first element on top without mutating the input.
• hasNext() is a loop, not a single check. One pop might reveal another nested list, which needs to be exploded again. The loop handles arbitrarily deep nesting.
• next() delegates to hasNext() first. This means next() is safe to call even if you forgot to check hasNext() beforehand, though the interface contract still expects callers to check.

Step-by-Step Walkthrough

Here is how the stack evolves as you call next() repeatedly on [[1,1],2,[1,1]]:

Notice that the stack never holds the fully expanded tree. It only expands one nested list at a time. When you call next() for the third time (to get the 2), the second inner list [1,1] is still sitting on the stack unexploded. It only gets expanded when you actually need those elements.

Complexity

Each element is pushed onto and popped from the stack at most twice (once during initialization or explosion, once when consumed). The total work across all calls is O(n) where n is the total number of integers in the structure. Individual hasNext() calls can take O(d) time where d is the nesting depth, but amortized over all calls the total is O(n).

Building Blocks

Three ideas combine here:

1. Stack for depth-first traversal. A stack naturally mimics the call stack that a recursive DFS would use. Pushing elements in reverse order ensures the leftmost element is processed first.
2. Lazy evaluation. You do not flatten the structure upfront. You expand each nested list only when you reach it during iteration. This is the iterator pattern applied to trees.
3. Invariant maintenance. The invariant of hasNext() is: after it returns True, the top of the stack is always an integer. next() relies on this invariant. Keeping the invariant in hasNext() makes next() trivial.

If you have solved Implement Stack Using Queues or Implement Queue Using Stacks, you are already familiar with using one data structure to simulate another. Here you are using a stack to simulate recursive tree traversal.

Edge Cases

Empty nested list. If nestedList is [], the stack starts empty. hasNext() returns False immediately. next() would call hasNext() and then try to pop from an empty stack, but the LeetCode guarantee is that next() is only called when hasNext() is True.

Deeply nested single integer. Something like [[[[[42]]]]]. Each hasNext() call explodes one layer of nesting. After 5 calls to the while loop body, the stack contains just the integer 42. Then hasNext() returns True and next() returns 42.

Empty inner lists. [[],1,[],[],2] has inner lists with no elements. When you explode an empty list, reversed([]) produces nothing, so nothing is pushed. The while loop continues to the next item. Empty lists are silently skipped.

Mixed nesting levels. [1,[2,[3,4],5],6] has integers and lists at the same nesting level. The stack handles this without any special casing: it just checks isInteger() at the top and either returns or explodes.

From Understanding to Recall

The pattern worth locking in is this sequence:

Push reversed. Peek. If list, explode. Repeat. If integer, yield.

When you see an iterator design problem in an interview, reach for this template. The stack replaces recursion, reversed insertion replaces DFS ordering, and the loop in hasNext() replaces the recursive call stack.

To make this stick with spaced repetition, practice writing hasNext() from scratch. That is the hard part. __init__ and next() are just one-liners once you have hasNext(). During each practice session, ask yourself: why do we reverse before pushing? What happens if there is a deeply nested empty list? Getting those answers right in under 30 seconds means the solution is fully internalized.

Related Posts

• Implement Stack Using Queues -- simulating a stack with queues, another design problem that tests how well you understand the iterator interface.
• Implement Queue Using Stacks -- the reverse simulation, same category of problem.
• Binary Search Tree Iterator -- another lazy iterator that uses a stack to traverse a tree one node at a time.