Design a flattening iterator over nested lists

01The design, in the six-step framework

1. Requirements and use cases

Two prompts share one design: flatten a 2D vector, and flatten an arbitrarily nested list. The interviewer is looking for laziness. A solution that flattens everything in the constructor passes the tests and fails the discussion.

• A caller builds the iterator from a nested list where each item is an integer or another nested list, to any depth.
• A caller calls next and gets integers in left-to-right depth-first order.
• A caller calls has_next, which must be true only if an integer really remains.
• Empty lists at any depth are skipped silently.
• Follow-up: a remove on the 2D vector deletes the element last returned.

Out of scope: mutation of the input while iterating, cycles in the input (mentioned below), and reverse iteration. Non-functional constraints: construction O(1), no copying of the input; memory O(depth), not O(total elements); and amortised O(1) per next.

2. Core entities and relationships

• NestedItem: the element interface; either a Leaf holding an integer or a Branch holding children.
• FlattenIterator: the public object; owns a stack of child cursors.
• Cursor: a Python iterator over one list's children; the stack holds one per open level.

The data is a Composite: a leaf and a branch answer the same question, so the walker never casts or checks a concrete type. The iterator aggregates the root it was given and composes its stack; it owns no element, only positions in the caller's structure. That stack is the explicit form of the recursion a depth-first walk would otherwise use, which is what lets the walk pause between calls.

3. Interfaces and key methods

from typing import Iterator, List, Optional, Protocol


class NestedItem(Protocol):
    def is_leaf(self) -> bool: """True for an integer, False for a list."""
    def value(self) -> int: """The integer; only valid on a leaf."""
    def children(self) -> List["NestedItem"]: """Direct children; only valid on a branch."""


class FlattenIterator:
    def __init__(self, root: List[NestedItem]) -> None:
        """O(1): pushes one cursor over the root and reads nothing."""

    def has_next(self) -> bool:
        """Advance to the next integer if needed; idempotent between next calls."""

    def __next__(self) -> int:
        """Return that integer and clear it; raises StopIteration when none remain."""

Invariants: the stack holds one cursor for each open list from the root down to the current position; ready holds an integer found but not yet returned, or nothing; and calling has_next repeatedly does not skip anything, which is why the found integer is parked rather than returned.

4. Design patterns and the choices behind them

An explicit stack of iterators

The force is pausing a depth-first walk between calls. Recursion keeps its position on the call stack, which disappears when the function returns, so a lazy walk has to hold that position itself. A stack of iterators is exactly that, and its size is the current depth. Python generators do this for free with yield from, and the code shows it as the short version; in an interview, offer it, then write the stack, because the question is about the mechanism. The rejected eager version copies everything into a flat list, costs O(n) up front, and is pointless when a caller stops early.

has_next does the work

The force is empty lists. Input [[], [[]]] contains no integers, so has_next must be false, and only walking can discover that. Putting the search in has_next, parking the result, and having next just hand it over means one code path decides what comes next. Designing it the other way round, with next doing the search and has_next checking stack depth, is the classic wrong answer. This is a contract decision of the kind in method design.

Composite over type checks

Asking items is_leaf keeps the walker independent of concrete classes; the 2D vector is simply a list of branches holding only leaves. Using encapsulation this way lets the same iterator serve both prompts.

5. Core flow and code

Flow for [1, [], [[2]]]:

1. The constructor pushes a cursor over the root.
2. has_next pulls 1, a leaf, and parks it. next returns it.
3. has_next pulls the empty list and pushes its cursor; that cursor is exhausted at once, so it is popped.
4. It pulls [[2]] and pushes, pulls [2] and pushes, pulls 2 and parks it.
5. After 2 is returned, each exhausted cursor pops in turn; the stack empties and has_next is false.

from dataclasses import dataclass
from typing import Iterator, List, Optional, Sequence, Union


class NestedItem:
    def is_leaf(self) -> bool:
        raise NotImplementedError

    def value(self) -> int:
        raise TypeError("not a leaf")

    def children(self) -> List["NestedItem"]:
        raise TypeError("not a branch")


@dataclass(frozen=True)
class Leaf(NestedItem):
    number: int

    def is_leaf(self) -> bool:
        return True

    def value(self) -> int:
        return self.number


@dataclass(frozen=True)
class Branch(NestedItem):
    items: tuple

    def is_leaf(self) -> bool:
        return False

    def children(self) -> List[NestedItem]:
        return list(self.items)


def build(raw: Sequence[Union[int, list]]) -> List[NestedItem]:
    """Test helper: Python lists and ints into the item model."""
    return [Leaf(x) if isinstance(x, int) else Branch(tuple(build(x))) for x in raw]


class FlattenIterator:
    def __init__(self, root: List[NestedItem]) -> None:
        self.stack: List[Iterator[NestedItem]] = [iter(root)]
        self.ready: Optional[int] = None
        self.has_ready = False

    def has_next(self) -> bool:
        while not self.has_ready and self.stack:
            item = next(self.stack[-1], None)
            if item is None:
                self.stack.pop()
            elif item.is_leaf():
                self.ready, self.has_ready = item.value(), True
            else:
                self.stack.append(iter(item.children()))
        return self.has_ready

    def __next__(self) -> int:
        if not self.has_next():
            raise StopIteration
        self.has_ready = False
        return self.ready  # type: ignore[return-value]

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


def flatten_gen(items: List[NestedItem]) -> Iterator[int]:
    """The generator version: the runtime keeps the stack for us."""
    for item in items:
        if item.is_leaf():
            yield item.value()
        else:
            yield from flatten_gen(item.children())


class Vector2D:
    """The 2D prompt plus remove(): indexes instead of iterators, so deletion is possible."""

    def __init__(self, rows: List[List[int]]) -> None:
        self.rows, self.r, self.c = rows, 0, 0
        self.last: Optional[tuple] = None

    def has_next(self) -> bool:
        while self.r < len(self.rows) and self.c >= len(self.rows[self.r]):
            self.r, self.c = self.r + 1, 0
        return self.r < len(self.rows)

    def __next__(self) -> int:
        if not self.has_next():
            raise StopIteration
        self.last = (self.r, self.c)
        self.c += 1
        return self.rows[self.r][self.c - 1]

    def remove(self) -> None:
        if self.last is None:
            raise RuntimeError("next must be called before remove")
        r, c = self.last
        del self.rows[r][c]
        if r == self.r:
            self.c -= 1
        self.last = None


if __name__ == "__main__":
    data = build([1, [], [[2]], [[]], [3, [4]]])
    assert list(FlattenIterator(data)) == [1, 2, 3, 4] == list(flatten_gen(data))
    it = FlattenIterator(build([[], [[]]]))
    assert not it.has_next() and not it.has_next()
    v = Vector2D([[1, 2], [], [3]])
    assert next(v) == 1
    v.remove()
    assert next(v) == 2 and next(v) == 3 and not v.has_next()

6. Extensions, concurrency and what interviewers probe

• Empty lists. [[], [[]]] is the test; a has_next based on stack depth says true and then fails.
• Lazy versus eager. Say the memory numbers: O(depth) against O(n), and that eager construction also pays for elements nobody reads.
• Cycles. A list that contains itself would walk forever; tracking the ids of open lists detects it and costs O(depth) memory. Raise a clear error rather than looping.
• Remove on the 2D vector. Vector2D uses indexes because a Python list iterator cannot delete; the cursor shifts back when the deleted slot was in the current row.
• Mutation during iteration. Unsupported; say so. A version counter on the input would fail fast, as standard collections do.
• What to test first. A flat list, a deep single chain, and the all-empty case; then call has_next three times before each next to prove it does not skip.
• Common mistake. Testing truthiness of the value, so a 0 is skipped. The code tests the item against None and uses a separate has_ready flag for the parked integer.

02Where this question was reported

• LeetCode Discuss · LeetCode 251 (Medium), Google-tagged, freq 18.2
• techprep.app · TechPrep LLD '2D Vector Class'
• LeetCode Discuss · LeetCode 341 (Medium), Google-tagged, freq 29.1