LeetCode 281 - Zigzag Iterator

LeetCode Problem 281

Difficulty: 🟡 Medium
Topics: Array, Design, Queue, Iterator

Solution

Problem Understanding

The problem asks us to design an iterator that alternates between two input arrays, returning elements in a cyclic or zigzag order. Instead of consuming all elements from the first vector and then moving to the second vector, we must switch back and forth between them whenever possible.

For example, if the input vectors are [1, 2] and [3, 4, 5, 6], the iterator should produce values in this order:

1, 3, 2, 4, 5, 6

The iterator begins with the first vector, then takes one element from the second vector, then returns to the first vector, and so on. Once one vector is exhausted, the iterator continues consuming the remaining elements from the other vector.

The class exposes three operations:

• The constructor initializes the iterator with two vectors.
• next() returns the next element in zigzag order.
• hasNext() reports whether any elements remain.

The constraints are relatively small, with at most 2000 total elements. This means almost any reasonable linear solution will pass comfortably. However, the problem is fundamentally about iterator design and extensibility, especially for the follow-up involving k vectors.

An important detail is that the iterator must behave incrementally. We are not simply asked to construct the final output array in advance. Instead, the iterator should maintain internal state between calls to next().

Several edge cases are important:

• One vector may be empty.
• Both vectors may have very different lengths.
• One vector may become exhausted earlier than the other.
• The iterator must never attempt to access an out-of-bounds index.
• hasNext() and next() must remain consistent throughout iteration.

The follow-up generalizes the problem from two vectors to k vectors. This strongly hints that we should think in terms of a queue of active iterators or indices rather than hardcoding alternation between exactly two arrays.

Approaches

Brute Force Approach

A straightforward approach is to precompute the entire zigzag traversal during construction. We can maintain two indices, repeatedly append elements alternately into a result array, and then use another pointer to iterate through that precomputed array.

This works because the zigzag order is deterministic. Once the merged order is built, each next() operation simply returns the next element from the precomputed list.

However, this approach is not ideal from a design perspective. It eagerly computes the entire traversal even if only a few elements are requested. It also does not extend naturally to the follow-up involving k vectors, because the logic becomes increasingly cumbersome as the number of vectors grows.

Optimal Approach

A better solution is to model the problem as a queue of active iterators.

The key insight is that every vector that still contains remaining elements should participate in the cycle. After consuming one element from a vector, if that vector still has more elements, it should be placed back into the queue so it can participate again later.

This creates a natural round-robin traversal pattern.

For each active vector, we only need:

• A reference to the vector itself
• The current index inside that vector

When next() is called:

1. Remove the front iterator state from the queue.
2. Return its current element.
3. Advance its index.
4. If elements remain, push it back into the queue.

This approach is elegant, efficient, and directly generalizes to the k vector follow-up.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n) preprocessing, O(1) per query	O(n)	Precomputes entire traversal order
Optimal	O(1) amortized per operation	O(k)	Queue-based cyclic iterator, easily extends to k vectors

Here, n is the total number of elements and k is the number of active vectors.

Algorithm Walkthrough

Optimal Queue-Based Algorithm

1. Initialize a queue that stores active iterator states.

Each state contains:

• The vector itself
• The current index within that vector

We only add vectors that are non-empty. This avoids unnecessary checks later. 2. During construction, examine both vectors.

If v1 contains elements, add (v1, 0) to the queue.

If v2 contains elements, add (v2, 0) to the queue.

The index 0 represents the next element to return from that vector. 3. When next() is called, remove the front state from the queue.

This represents the vector whose turn it currently is. 4. Read the element at the stored index.

This is the next zigzag element to return. 5. Advance the index by one.

After consuming an element, the iterator state must move forward. 6. If the vector still has remaining elements, push the updated state back into the queue.

This preserves cyclic ordering. The vector will participate again after the other active vectors have had their turn. 7. Return the extracted element. 8. For hasNext(), simply check whether the queue is non-empty.

If the queue still contains iterator states, then at least one vector still has remaining elements.

Why it works

The queue always maintains the invariant that its elements represent vectors with remaining unconsumed values, ordered by whose turn comes next in the cyclic traversal.

Every call to next() consumes exactly one element from the front vector and then re-enqueues that vector only if additional elements remain. This guarantees that:

• Elements are returned in cyclic order.
• No element is skipped.
• No vector participates after being exhausted.

Because every vector rotates through the queue fairly, the traversal exactly matches the required zigzag behavior.

Python Solution

from collections import deque
from typing import List, Deque, Tuple

class ZigzagIterator:
    def __init__(self, v1: List[int], v2: List[int]):
        self.queue: Deque[Tuple[List[int], int]] = 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

# Your ZigzagIterator object will be instantiated and called as such:
# i, v = ZigzagIterator(v1, v2), []
# while i.hasNext(): v.append(i.next())

The implementation uses a deque because we repeatedly remove elements from the front and append elements to the back. A deque provides efficient O(1) operations for both.

The constructor initializes the queue with only the non-empty vectors. Each queue entry stores both the vector and the next index to consume from that vector.

The next() method removes the front iterator state, retrieves the current value, advances the index, and re-adds the vector to the queue if it still contains elements.

The hasNext() method is simple because the queue itself represents all remaining active iterators. If the queue is empty, iteration is complete.

This design naturally extends to the follow-up involving k vectors. Instead of handling two vectors specially, we can initialize the queue with all non-empty vectors.

Go Solution

type IteratorState struct {
	vector []int
	index  int
}

type ZigzagIterator struct {
	queue []IteratorState
}

func Constructor(v1, v2 []int) *ZigzagIterator {
	z := &ZigzagIterator{
		queue: []IteratorState{},
	}

	if len(v1) > 0 {
		z.queue = append(z.queue, IteratorState{
			vector: v1,
			index:  0,
		})
	}

	if len(v2) > 0 {
		z.queue = append(z.queue, IteratorState{
			vector: v2,
			index:  0,
		})
	}

	return z
}

func (this *ZigzagIterator) next() int {
	current := this.queue[0]
	this.queue = this.queue[1:]

	value := current.vector[current.index]
	current.index++

	if current.index < len(current.vector) {
		this.queue = append(this.queue, current)
	}

	return value
}

func (this *ZigzagIterator) hasNext() bool {
	return len(this.queue) > 0
}

/**
 * Your ZigzagIterator object will be instantiated and called as such:
 * obj := Constructor(param_1, param_2);
 * for obj.hasNext() {
 *	 ans = append(ans, obj.next())
 * }
 */

The Go implementation follows the same queue-based design as the Python version. Since Go does not have a built-in deque structure, we simulate the queue using slices.

Removing the front element is done using:

this.queue = this.queue[1:]

Appending back into the queue uses Go's built-in append() function.

Unlike Python, Go distinguishes between nil slices and empty slices, but both behave correctly with len(slice) == 0, so no special handling is required.

Integer overflow is not a concern because the iterator only stores indices within arrays of size at most 1000.

Worked Examples

Example 1

Input:

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

Initial queue state:

Queue Front → Back
([1,2], 0), ([3,4,5,6], 0)

Step-by-step execution:

Step	Popped State	Returned Value	Reinserted State	Queue After
1	([1,2], 0)	1	([1,2], 1)	([3,4,5,6],0), ([1,2],1)
2	([3,4,5,6],0)	3	([3,4,5,6],1)	([1,2],1), ([3,4,5,6],1)
3	([1,2],1)	2	none	([3,4,5,6],1)
4	([3,4,5,6],1)	4	([3,4,5,6],2)	([3,4,5,6],2)
5	([3,4,5,6],2)	5	([3,4,5,6],3)	([3,4,5,6],3)
6	([3,4,5,6],3)	6	none	empty

Final output:

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

Example 2

Input:

v1 = [1]
v2 = []

Initial queue:

Queue
([1],0)

Execution:

Step	Returned Value	Queue After
1	1	empty

Final output:

[1]

Example 3

Input:

v1 = []
v2 = [1]

Initial queue:

Queue
([1],0)

Execution:

| Step | Returned Value | Queue After |

|---|---|

| 1 | 1 | empty |

Final output:

[1]

Complexity Analysis

Measure	Complexity	Explanation
Time	O(1) amortized per operation	Each element is inserted and removed from the queue exactly once
Space	O(k)	The queue stores at most one active state per vector

The total work across the entire traversal is linear in the number of elements because every element is processed exactly once. Queue operations are constant time.

For the original problem with two vectors, k = 2, so the extra space usage is effectively constant. The formulation becomes more meaningful in the follow-up where multiple vectors are involved.

Test Cases

from typing import List

def collect(iterator) -> List[int]:
    result = []

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

    return result

# Example 1
it = ZigzagIterator([1, 2], [3, 4, 5, 6])
assert collect(it) == [1, 3, 2, 4, 5, 6]  # standard alternating behavior

# Example 2
it = ZigzagIterator([1], [])
assert collect(it) == [1]  # second vector empty

# Example 3
it = ZigzagIterator([], [1])
assert collect(it) == [1]  # first vector empty

# Both vectors same size
it = ZigzagIterator([1, 2, 3], [4, 5, 6])
assert collect(it) == [1, 4, 2, 5, 3, 6]  # perfect alternation

# First vector longer
it = ZigzagIterator([1, 2, 3, 4], [5, 6])
assert collect(it) == [1, 5, 2, 6, 3, 4]  # remaining tail from first vector

# Second vector longer
it = ZigzagIterator([1], [2, 3, 4, 5])
assert collect(it) == [1, 2, 3, 4, 5]  # remaining tail from second vector

# Single element each
it = ZigzagIterator([10], [20])
assert collect(it) == [10, 20]  # minimal balanced case

# Negative numbers
it = ZigzagIterator([-1, -2], [-3, -4])
assert collect(it) == [-1, -3, -2, -4]  # handles negatives correctly

# Large imbalance
it = ZigzagIterator([1], [2, 3, 4, 5, 6, 7])
assert collect(it) == [1, 2, 3, 4, 5, 6, 7]  # one vector exhausted immediately

# hasNext behavior
it = ZigzagIterator([1], [])
assert it.hasNext() is True
assert it.next() == 1
assert it.hasNext() is False  # iterator exhausted correctly

Test	Why
[1,2] and [3,4,5,6]	Validates standard zigzag traversal
[1] and []	Ensures empty second vector works
[] and [1]	Ensures empty first vector works
Equal-length vectors	Confirms perfect alternation
First vector longer	Verifies leftover handling
Second vector longer	Verifies asymmetric exhaustion
Single element vectors	Tests smallest balanced case
Negative values	Confirms value sign does not matter
Large imbalance	Tests immediate exhaustion of one vector
Explicit hasNext() checks	Ensures iterator state remains consistent

Edge Cases

One important edge case occurs when one vector is empty from the start. A naive implementation that blindly alternates between vectors may attempt to access an invalid index or repeatedly check an already exhausted vector. This implementation avoids the issue entirely by only inserting non-empty vectors into the queue during initialization.

Another important edge case is when one vector becomes exhausted before the other. For example, [1, 2] and [3, 4, 5, 6, 7]. Some implementations incorrectly continue alternating and accidentally skip elements or produce invalid accesses. In this solution, once a vector runs out of elements, it is simply not reinserted into the queue. The remaining vectors continue normally.

A subtle edge case involves repeated calls to hasNext() near the end of iteration. Some buggy iterator implementations accidentally modify internal state during hasNext(), causing inconsistencies between hasNext() and next(). This implementation keeps hasNext() completely passive by only checking whether the queue is empty. No internal state changes occur.

Another edge case involves vectors containing only one element each. The iterator must still alternate correctly and terminate cleanly. Because each vector is removed from the queue after its single element is consumed, the queue becomes empty immediately after both elements are returned.