Problem 79 · Object & System Design

Flatten Nested List Iterator

MediumClassic DSA & systems

01Problem

Given an arbitrarily nested list of integers, implement an iterator with next() and has_next() that yields the integers in order, lazily.

02Key insight

Keep a stack of iterators instead of flattening up front: has_next advances until the top yields an integer (pushing sub-iterators for lists, popping exhausted ones). Memory is O(depth), and empty sublists are handled naturally.

03Solution

solution.pypython36 lines

from typing import Iterator, Union

Nested = list[Union[int, "Nested"]]


class NestedIterator:
    def __init__(self, nested: Nested) -> None:
        self.stack: list[Iterator] = [iter(nested)]
        self.peeked: int | None = None
        self.has_peeked = False

    def has_next(self) -> bool:
        while not self.has_peeked and self.stack:
            item = next(self.stack[-1], _END)
            if item is _END:
                self.stack.pop()
            elif isinstance(item, list):
                self.stack.append(iter(item))
            else:
                self.peeked, self.has_peeked = item, True
        return self.has_peeked

    def next(self) -> int:
        if not self.has_next():
            raise StopIteration
        self.has_peeked = False
        return self.peeked  # type: ignore[return-value]

    def __iter__(self) -> "NestedIterator":
        return self

    def __next__(self) -> int:
        return self.next()


_END = object()

Try it first.

TimeO(1) amortized per element

SpaceO(depth)

tests.py · 4 passing

tests.pypython5 lines

assert list(NestedIterator([[1, 1], 2, [1, 1]])) == [1, 1, 2, 1, 1]
assert list(NestedIterator([1, [4, [6]]])) == [1, 4, 6]
assert list(NestedIterator([[], [[]], [[], [7]], []])) == [7]
it = NestedIterator([])
assert not it.has_next()

04Staff-level follow-up

Same pattern as streaming a JSON/Parquet nested column or walking a B-tree: a stack of cursors keeps memory at O(depth) so it works on inputs larger than RAM.