Design a zigzag (cyclic) iterator over k lists

01The design, in the six-step framework

1. Requirements and use cases

The two-list version is a warm-up. The real question comes straight after: "what if there are k lists", at which point the name changes to cyclic iterator and the design has to stop hard-coding two cursors.

• A caller builds the iterator from k iterators of any length, including empty ones.
• A caller calls next and receives elements in round-robin order: first of each, then second of each.
• When one source runs out, the rotation continues over the others in the same relative order.
• A caller asks has_next at any time; it is true while any source has an element left.
• Sources are lazy: the iterator never reads ahead more than one element per source.

Out of scope: concurrent consumers, weighted rotation (mentioned as an extension) and rewinding. Non-functional constraints: O(1) per next regardless of k; memory O(k) and independent of the total number of elements; works for infinite sources; and an exhausted source must leave the rotation permanently, so it costs nothing afterwards.

2. Core entities and relationships

• ZigzagIterator: the public object; owns the rotation queue and exposes the iterator protocol.
• Rotation queue: a deque of live sources; the front source supplies the next element.
• Source: a peeking iterator around each input, which lets the queue hold only sources known to have an element.

The zigzag iterator composes its queue and aggregates the sources it was handed. Wrapping each input in a peeking iterator is what makes the rule simple: anything in the queue is guaranteed to have a next element, so has_next becomes "is the queue non-empty", with no scanning. The shape is the fairness half of producer-consumer: several producers, one consumer, taken in turn.

3. Interfaces and key methods

from typing import Generic, Iterable, Iterator, TypeVar

T = TypeVar("T")


class ZigzagIterator(Generic[T]):
    def __init__(self, sources: Iterable[Iterator[T]]) -> None:
        """Keeps only the sources that have at least one element, in the given order."""

    def has_next(self) -> bool:
        """True while any source has an element left; O(1)."""

    def __next__(self) -> T:
        """Element from the front source; the source rejoins the back if it has more."""

    def __iter__(self) -> "ZigzagIterator[T]": ...

Invariants: every source in the queue has at least one element ready; the relative order of live sources never changes; and each input element is returned exactly once. __next__ raises StopIteration when the queue is empty, following the language's protocol as method design asks.

4. Design patterns and the choices behind them

A queue of sources, not an index and a skip loop

The force is sources running out at different times. The obvious version keeps an index and advances it, skipping exhausted sources; when only one of k sources remains, every call walks past k − 1 dead entries, so next becomes O(k). A deque removes a source the moment it is exhausted, which keeps every call O(1). This is the main point of the iterator design here: state lives in the structure, not in a counter with special cases.

Peek to decide re-queueing

The force is that an iterator cannot say "I am finished" without being read. After taking an element, the source must be checked before it goes to the back, or the queue would hold an empty source and has_next would lie. A one-element buffer per source answers that with bounded look-ahead. The rejected alternative catches StopIteration inside next and retries with the following source; it works, but has_next then has to scan, and it would read ahead by an unbounded number of sources.

Generic over iterators, not lists

Accepting iterators rather than lists means infinite sources and generators work, and it hides how each input is stored behind one small interface. That is encapsulation: the rotation never indexes into anything.

5. Core flow and code

Flow over sources [1, 2, 3], [], [4]:

1. The constructor wraps all three and drops the empty one; the queue holds A and C.
2. next pops A, takes 1; A has more, so it goes to the back. Queue: C, A.
3. next pops C, takes 4; C is finished and is dropped. Queue: A.
4. The next two calls return 2 and 3 from A, after which the queue is empty and has_next is false.

from collections import deque
from typing import Deque, Generic, Iterable, Iterator, List, Optional, Tuple, TypeVar

T = TypeVar("T")


class PeekingSource(Generic[T]):
    """One-element look-ahead around an input iterator."""

    def __init__(self, source: Iterator[T]) -> None:
        self.source = source
        self.buffered = False
        self.buffer: Optional[T] = None

    def has_next(self) -> bool:
        if self.buffered:
            return True
        try:
            self.buffer = next(self.source)
        except StopIteration:
            return False
        self.buffered = True
        return True

    def take(self) -> T:
        if not self.has_next():
            raise StopIteration
        self.buffered = False
        value, self.buffer = self.buffer, None
        return value  # type: ignore[return-value]


class ZigzagIterator(Generic[T]):
    def __init__(self, sources: Iterable[Iterator[T]]) -> None:
        self.queue: Deque[PeekingSource[T]] = deque()
        for source in sources:
            wrapped = PeekingSource(source)
            if wrapped.has_next():
                self.queue.append(wrapped)

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

    def __next__(self) -> T:
        if not self.queue:
            raise StopIteration
        source = self.queue.popleft()
        value = source.take()
        if source.has_next():
            self.queue.append(source)
        return value

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


class WeightedZigzag(ZigzagIterator[T]):
    """Extension: take `burst` elements from a source before rotating (burst=1 is plain zigzag)."""

    def __init__(self, sources: Iterable[Iterator[T]], burst: int = 2) -> None:
        super().__init__(sources)
        self.burst = burst
        self.taken = 0

    def __next__(self) -> T:
        if not self.queue:
            raise StopIteration
        source = self.queue[0]
        value = source.take()
        self.taken += 1
        if not source.has_next():
            self.queue.popleft()
            self.taken = 0
        elif self.taken == self.burst:
            self.queue.rotate(-1)
            self.taken = 0
        return value


if __name__ == "__main__":
    z = ZigzagIterator([iter([1, 2, 3]), iter([]), iter([4])])
    assert list(z) == [1, 4, 2, 3] and not z.has_next()
    y = ZigzagIterator([iter("ab"), iter("cdef"), iter("g")])
    assert "".join(y) == "acgbdef"
    w = WeightedZigzag([iter([1, 2, 3]), iter([4, 5, 6])], burst=2)
    assert list(w) == [1, 2, 4, 5, 3, 6]
    e: ZigzagIterator[int] = ZigzagIterator([])
    assert not e.has_next()

6. Extensions, concurrency and what interviewers probe

• From two to k. The queue version needs no change; a two-pointer version needs rewriting. Starting with the general shape when the prompt says two is a fair call if you say why.
• Cost when sources finish unevenly. The skip-loop version degrades to O(k) per call. Quote that number; it is the reason for the deque.
• Weighted rotation. WeightedZigzag takes a burst per source; per-source weights would store a count beside each entry.
• Merging by value instead of by turn. If the follow-up orders elements, the deque becomes a heap keyed by each source's next element; the wrapper stays the same.
• Thread safety. One consumer is the contract. For several, a lock around __next__ is enough, because pop, take and re-append form one critical section.
• What to test first. Empty sources in the first, middle and last position, then one source much longer than the others.
• Common mistake. Re-queueing before checking for more, so has_next returns true and the following next raises.

02Where this question was reported

• LeetCode Discuss · LeetCode 281 (Medium), Google-tagged, freq 16.6
• evelynn.gitbooks.io · Google-tagged study collection; test peek-before-next and next-before-peek; generics follow-up
• evelynn.gitbooks.io · Zigzag iterator; 'cyclic' once >2 lists; queue of sub-iterators
• Taro · Taro claims Google/Meta
• Medium · 'Real Google' peeking iterator write-up