Design independent level-order iterators over a shared tree

01The design, in the six-step framework

1. Requirements and use cases

A new-grad phone screen that starts as tree traversal and turns into a design conversation: an iterator returns one level per call, several iterators walk one tree at once, and then the interviewer asks what happens if the tree changes.

• A caller creates a level iterator over a binary tree.
• Each next returns the values of one level, left to right, top level first.
• A caller checks has_next before calling next.
• Several independent iterators run over the same tree at their own pace without affecting one another.
• Follow-up: a writer mutates the tree while iterators are open.

Out of scope: traversal orders other than level order, persistence, and balancing. Non-functional constraints: the walk is lazy, doing work for a level only when it is requested; memory per iterator is O(width of one level), not O(n); the tree itself holds no iterator state, because that is what makes iterators independent.

2. Core entities and relationships

• TreeNode: a value and two children; read-only from an iterator's point of view.
• LevelIterator: one walk; holds only the nodes of the next level.
• VersionedTree: the follow-up; the root plus a version counter and a read-write lock, supporting mutation.
• SnapshotIterator: pins one immutable root version, so concurrent writes cannot disturb it.

An iterator aggregates the tree and owns only its frontier list. Nodes know nothing about iterators; if a visited flag lived on the node, two iterators would corrupt each other, so all walking state stays in the iterator. That single rule delivers the independence requirement with no locks at all while the tree is read-only.

3. Interfaces and key methods

from typing import Generic, List, Optional, TypeVar

T = TypeVar("T")


class TreeNode(Generic[T]):
    value: T
    left: "Optional[TreeNode[T]]"
    right: "Optional[TreeNode[T]]"


class LevelIterator(Generic[T]):
    def __init__(self, root: Optional[TreeNode[T]]) -> None:
        """O(1): the frontier starts as just the root; an empty tree has no levels."""

    def has_next(self) -> bool:
        """True while the frontier is non-empty. Pure: does no walking."""

    def __next__(self) -> List[T]:
        """Values of the frontier; builds the following frontier. O(width of two levels)."""


class VersionedTree(Generic[T]):
    def levels(self) -> LevelIterator[T]:
        """An iterator pinned to the current immutable root."""

    def insert(self, path: str, value: T) -> None:
        """path of 'L'/'R' to the new node; copies the nodes on that path, then publishes."""

Invariants: the frontier holds exactly the nodes at the next unreturned level, left to right; a pinned root and everything under it are never mutated; and publishing a new root is one reference assignment.

4. Design patterns and the choices behind them

A frontier list instead of a queue

The force is the memory claim. A usual BFS queue holds parts of two levels and needs a size counter to find the boundary. Because next returns exactly one level, the iterator can keep a plain list for the current level and build a new list for the next one. Peak memory is two adjacent levels, which is O(width). In a full binary tree the last level holds about half the nodes, so width is still O(n) in the worst case; saying this stops you overselling it, and a chain has width 1.

Copy-on-write for the mutation follow-up

The force is an iterator that has half a tree in its frontier when a writer arrives. There are three answers. A read-write lock held for each next gives each level consistency but not the whole walk: level 3 may reflect a write that level 2 did not. Holding a read lock for the whole walk blocks writers for as long as the slowest consumer dawdles, which is unbounded. A fail-fast version counter raises an error in the iterator when the tree changed, which is honest and cheap, and is what standard collections do. The chosen answer is a persistent snapshot: the writer copies the path from root to the change, sharing every other subtree, and publishes a new root. Open iterators keep the old root and see a consistent tree forever, without locks on the read path.

A pure has_next

Here the frontier already answers the question, so has_next does no work and can be called freely. Stating which calls have side effects is part of method design.

5. Core flow and code

Flow for insert("R", 9) while an iterator holds version 1 of a tree with root 1 and left child 2:

1. The writer takes the write lock, which serialises writers only.
2. It builds a copy of the root with right child 9, sharing the node 2 unchanged.
3. It assigns the new root and raises the version. That assignment is the publish.
4. The open iterator still holds the old root, so its levels are [1], [2]. A new iterator sees [1], [2, 9].

import threading
from dataclasses import dataclass, replace
from typing import Generic, Iterator, List, Optional, Sequence, TypeVar

T = TypeVar("T")


@dataclass(frozen=True)
class TreeNode(Generic[T]):
    value: T
    left: "Optional[TreeNode[T]]" = None
    right: "Optional[TreeNode[T]]" = None


class LevelIterator(Generic[T]):
    def __init__(self, root: Optional[TreeNode[T]]) -> None:
        self.frontier: List[TreeNode[T]] = [root] if root is not None else []

    def has_next(self) -> bool:
        return bool(self.frontier)

    def __next__(self) -> List[T]:
        if not self.frontier:
            raise StopIteration
        values = [node.value for node in self.frontier]
        following: List[TreeNode[T]] = []
        for node in self.frontier:
            if node.left is not None:
                following.append(node.left)
            if node.right is not None:
                following.append(node.right)
        self.frontier = following
        return values

    def __iter__(self) -> "LevelIterator[T]":
        return self


class VersionedTree(Generic[T]):
    """Copy-on-write tree: nodes are frozen, so a pinned root never changes under a reader."""

    def __init__(self, root: Optional[TreeNode[T]] = None) -> None:
        self.root = root
        self.version = 0
        self.write_lock = threading.Lock()

    def levels(self) -> LevelIterator[T]:
        return LevelIterator(self.root)          # reading one reference is atomic

    def insert(self, path: str, value: T) -> None:
        with self.write_lock:
            self.root = self._with(self.root, path, value)
            self.version += 1

    def _with(self, node: Optional[TreeNode[T]], path: str, value: T) -> TreeNode[T]:
        if not path:
            if node is None:
                return TreeNode(value)
            return replace(node, value=value)
        if node is None:
            raise KeyError("parent on path is missing")
        if path[0] == "L":
            return replace(node, left=self._with(node.left, path[1:], value))
        if path[0] == "R":
            return replace(node, right=self._with(node.right, path[1:], value))
        raise ValueError("path uses only L and R")


if __name__ == "__main__":
    tree = TreeNode(1, TreeNode(2, TreeNode(4)), TreeNode(3, None, TreeNode(5)))
    a, b = LevelIterator(tree), LevelIterator(tree)
    assert next(a) == [1] and next(a) == [2, 3]
    assert next(b) == [1]                         # independent of a
    assert next(a) == [4, 5] and not a.has_next()
    assert list(b) == [[2, 3], [4, 5]]
    assert list(LevelIterator(None)) == []

    vt: VersionedTree[int] = VersionedTree(TreeNode(1, TreeNode(2)))
    pinned = vt.levels()
    assert next(pinned) == [1]
    vt.insert("R", 9)
    assert list(pinned) == [[2]]                  # old snapshot unaffected
    assert list(vt.levels()) == [[1], [2, 9]] and vt.version == 1

6. Extensions, concurrency and what interviewers probe

• Where iterator state lives. Only in the iterator. Say this first; it answers the independence question before it is asked.
• Memory. O(width), and name the worst case: a full tree has about n / 2 nodes in its last level.
• Is a lock needed for read-only sharing? No. Immutable data shared by readers needs no synchronisation, and adding a lock anyway is a mark against you.
• Publishing safely. In Python the reference assignment is atomic; in Java, make the root field volatile, or a reader may see a partly constructed node.
• Writer cost. O(height) new nodes per write. Old versions are reclaimed by the garbage collector once no iterator holds them; in a language without one, this needs reference counting.
• Common mistake. Using one shared queue on the tree object, or marking nodes visited, so that a second iterator silently gets nothing.

02Where this question was reported

• FastPrep · next() returns one level; memory scales with level width; follow-up: multiple iterators sharing a read-only tree, locking/version snapshots if mutation allowed