Problem 80 · Object & System Design

Peeking Iterator

MediumClassic DSA & systems

01Problem

Wrap any iterator to support peek() (return next element without advancing), next(), and has_next(). Must work for iterators that may yield None.

02Key insight

Cache exactly one lookahead element plus a separate 'has cached' flag (not a None check, since None is a valid element). This is the building block for k-way merges and merge joins.

03Solution

solution.pypython38 lines

from typing import Generic, Iterable, Iterator, TypeVar

T = TypeVar("T")


class PeekingIterator(Generic[T]):
    def __init__(self, source: Iterable[T]) -> None:
        self._it: Iterator[T] = iter(source)
        self._buffered = False
        self._head: T | None = None

    def _fill(self) -> bool:
        if not self._buffered:
            try:
                self._head = next(self._it)
                self._buffered = True
            except StopIteration:
                return False
        return True

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

    def peek(self) -> T:
        if not self._fill():
            raise StopIteration
        return self._head  # type: ignore[return-value]

    def next(self) -> T:
        value = self.peek()
        self._buffered, self._head = False, None
        return value

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

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

Try it first.

TimeO(1) per op

SpaceO(1)

tests.py · 4 passing

tests.pypython6 lines

p = PeekingIterator([1, 2, 3])
assert p.next() == 1 and p.peek() == 2 and p.peek() == 2 and p.next() == 2
assert p.next() == 3 and not p.has_next()
q = PeekingIterator([None, 0])
assert q.has_next() and q.peek() is None and q.next() is None and q.next() == 0
assert list(PeekingIterator(x for x in "ab")) == ["a", "b"]

04Staff-level follow-up

Use it for an external k-way merge of sorted runs: a heap of (peek(), run_id) across PeekingIterators streams the merge with O(k) memory, the core of LSM compaction and sort-merge joins.