LeetCode 281: Zigzag Iterator

Problem Restatement

We are given two integer lists, v1 and v2.

We need to implement an iterator that returns their elements alternately.

For example:

v1 = [1, 2]
v2 = [3, 4, 5, 6]

The iterator should return:

1, 3, 2, 4, 5, 6

Once one list runs out, we continue returning elements from the other list.

The class has three operations:

ZigzagIterator(v1, v2)
next()
hasNext()

The source statement defines next() as returning the current element and moving the iterator forward, and hasNext() as checking whether elements remain. It also asks how the design extends to k vectors. The follow-up says that for k > 2, the order can be treated as cyclic.

Input and Output

Class shape:

class ZigzagIterator:
    def __init__(self, v1: list[int], v2: list[int]):
        ...

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

    def hasNext(self) -> bool:
        ...

Examples

For:

v1 = [1, 2]
v2 = [3, 4, 5, 6]

Calls should behave like this:

So the full output is:

[1, 3, 2, 4, 5, 6]

For:

v1 = [1]
v2 = []

The output is:

[1]

For:

v1 = []
v2 = [1]

The output is:

[1]

First Thought: Build the Whole Answer

A simple solution is to build the entire zigzag output during initialization.

We can use two pointers, one for each list.

class ZigzagIterator:
    def __init__(self, v1: list[int], v2: list[int]):
        self.values = []
        i = 0

        while i < len(v1) or i < len(v2):
            if i < len(v1):
                self.values.append(v1[i])
            if i < len(v2):
                self.values.append(v2[i])
            i += 1

        self.index = 0

    def next(self) -> int:
        value = self.values[self.index]
        self.index += 1
        return value

    def hasNext(self) -> bool:
        return self.index < len(self.values)

This works for two lists.

For v1 = [1, 2] and v2 = [3, 4, 5, 6], it builds:

[1, 3, 2, 4, 5, 6]

Then next() simply returns one value at a time.

Problem With Prebuilding

Prebuilding stores all elements in a separate list.

That means extra space is proportional to the total number of elements.

Also, the follow-up asks about extending the design to k vectors.

For many lists, it is cleaner to keep track of active iterators instead of manually writing special logic for v1, v2, v3, and so on.

Key Insight

Use a queue of active iterators.

Each queue entry represents one list that still has values left.

When next() is called:

1. Remove the first iterator from the queue.
2. Return its next value.
3. If the iterator still has more values, put it back at the end of the queue.

This creates cyclic order naturally.

For two vectors:

v1, v2, v1, v2, ...

If v1 runs out, it disappears from the queue, and the iterator continues with v2.

This is also exactly the same idea needed for the k-vector follow-up.

Algorithm

During initialization:

1. Create an empty queue.
2. If v1 is non-empty, push [v1, 0].
3. If v2 is non-empty, push [v2, 0].

Each queue item stores:

[vector, index]

where index is the next unread position in that vector.

For next():

1. Pop the leftmost queue item.
2. Read vector[index].
3. Increment index.
4. If the vector still has more elements, append it back.
5. Return the value.

For hasNext():

return len(queue) > 0

Correctness

The queue always stores exactly the vectors that still have unread elements.

The front of the queue is the vector whose turn comes next.

When next() is called, the algorithm removes the front vector and returns its next unread value. This matches the required alternating order because vectors are processed in queue order.

If that vector still has unread elements, the algorithm appends it to the back of the queue. This gives every other active vector a turn before this vector is used again.

If that vector has no unread elements left, the algorithm does not reinsert it. Therefore exhausted vectors are skipped automatically.

So every call to next() returns the next value in cyclic order among all non-empty vectors. Since each returned value advances exactly one index and no index is repeated, every input element is returned exactly once.

Complexity

Let N = len(v1) + len(v2).

For the k-vector version, the constructor uses O(k) space for active iterators.

Implementation

from collections import deque

class ZigzagIterator:
    def __init__(self, v1: list[int], v2: list[int]):
        self.queue = deque()

        if v1:
            self.queue.append([v1, 0])

        if v2:
            self.queue.append([v2, 0])

    def next(self) -> int:
        vector, index = self.queue.popleft()

        value = vector[index]
        index += 1

        if index < len(vector):
            self.queue.append([vector, index])

        return value

    def hasNext(self) -> bool:
        return len(self.queue) > 0

Code Explanation

We use a queue because the iterator must rotate through the active vectors:

self.queue = deque()

During initialization, only non-empty vectors are added:

if v1:
    self.queue.append([v1, 0])

if v2:
    self.queue.append([v2, 0])

Each item stores the vector and the next index to read.

In next(), we take the vector whose turn comes next:

vector, index = self.queue.popleft()

Then we read its current value:

value = vector[index]

After reading, we advance the index:

index += 1

If the vector still has values left, we put it at the back of the queue:

if index < len(vector):
    self.queue.append([vector, index])

Finally, we return the value:

return value

hasNext() checks whether any active vector remains:

return len(self.queue) > 0

Extension to k Vectors

The same queue design extends directly to k vectors.

Instead of checking only v1 and v2, we loop over all vectors.

from collections import deque

class ZigzagIteratorK:
    def __init__(self, vectors: list[list[int]]):
        self.queue = deque()

        for vector in vectors:
            if vector:
                self.queue.append([vector, 0])

    def next(self) -> int:
        vector, index = self.queue.popleft()

        value = vector[index]
        index += 1

        if index < len(vector):
            self.queue.append([vector, index])

        return value

    def hasNext(self) -> bool:
        return len(self.queue) > 0

For:

vectors = [
    [1, 2, 3],
    [4, 5, 6, 7],
    [8, 9],
]

The output is:

[1, 4, 8, 2, 5, 9, 3, 6, 7]

Testing

def collect(iterator):
    result = []

    while iterator.hasNext():
        result.append(iterator.next())

    return result

def test_zigzag_iterator():
    assert collect(ZigzagIterator([1, 2], [3, 4, 5, 6])) == [
        1, 3, 2, 4, 5, 6
    ]

    assert collect(ZigzagIterator([1], [])) == [1]

    assert collect(ZigzagIterator([], [1])) == [1]

    assert collect(ZigzagIterator([], [1, 2, 3])) == [1, 2, 3]

    assert collect(ZigzagIterator([1, 2, 3], [4])) == [1, 4, 2, 3]

    print("all tests passed")

test_zigzag_iterator()

Test meaning: