Flatten Nested List Iterator

Implement an iterator that flattens arbitrarily nested lists in depth-first order.

iterator generator flatten nested lazy

Write a class `NestedIterator` that takes a nested list of integers (where an element is either an integer or a list of integers) and returns them in depth-first order.

Implement the following methods:
- `__init__(self, nested_list: list)` – stores the input and prepares the iterator.
- `__iter__(self)` – returns `self` (the iterator is its own iterable).
- `__next__(self)` – returns the next integer in depth-first order, or raises `StopIteration` when all integers have been yielded.

The input can contain lists at arbitrary depth, and empty lists can appear anywhere. The depth-first order means: for each element, if it is an integer, yield it; if it is a list, recursively traverse its elements in order.

Your implementation must be lazy: it should not flatten the entire list at construction time. Instead, it should traverse on demand as `__next__` is called. The class must work correctly with both `for` loops and manual `next()` calls.

You may use Python's generator machinery inside the iterator (e.g., a `yield` in a generator method) as long as the class satisfies the iterator protocol.

Implement the class exactly with the signature given.

Constraints

- The nested list can contain integers or lists, with any depth (including 0 depth, i.e., just an integer, and 1 depth, i.e., a list of integers).
- The total number of integers can be large (up to 10^5), but memory usage should be proportional to the nesting depth, not the total number of elements.
- All integers are within Python's int range.
- Each call to `__next__` must run in O(1) amortized time.

Example

>>> ni = NestedIterator([1, [2, [3, 4]], 5])
>>> list(ni)
[1, 2, 3, 4, 5]
>>> ni = NestedIterator([[1, 2], 3, [4, [5, [6]]]])
>>> iter(ni) is ni
True
>>> [next(ni) for _ in range(6)]
[1, 2, 3, 4, 5, 6]

10 points ~15 min

Recent Submissions

No submissions yet — hit Run Tests to try!

Hints

Think about how to recursively traverse a nested list lazily. A generator function can be used to yield integers in depth-first order, but you need to integrate it with `__next__`.

Store a generator object in `__init__` that does the traversal, then `__next__` can call `next` on that generator.

Remember that `yield from` is a clean way to recursively delegate to sublists.

Your iterator must raise `StopIteration` automatically when the generator is exhausted, so do not catch it unless you want to convert it.

Test Results

Press Ctrl+Enter or click Run Tests to execute your code.