Zigzag Iterator

How it works

Using a queue to manage pointers provides a cleaner and more efficient approach. Instead of scanning through all vectors on each call, we maintain a queue of (vector index, element index) pairs representing which elements are ready to be returned. After returning an element, we add the pointer for the next element in that vector (if any) to the back of the queue. This naturally handles vectors of different lengths and ensures O(1) time for each next() call.

Steps

1. Initialize a queue with pointers to the first element of each non-empty vector, stored as (vector index, element index) pairs.
2. In next():
   • Dequeue the front pointer to get the current vector and element indices.
   • Retrieve the element to return from the specified position.
   • If there are more elements in that vector, enqueue a pointer to the next element.
   • Return the retrieved element.
3. In hasNext(), simply check if the queue is non-empty.

Complexity

• Time Complexity: O(1)

  • For both the next() function and the hasNext() function, we have a constant time complexity, as we discussed before.

• Space Complexity: O(K)

  • We use a queue to keep track of the pointers to the input vectors in the variable self.vectors.

  • As a result, we would need O(K) space for K vectors.

  • Although the size of the queue will reduce over time once we exhaust some shorter vectors, the space complexity for both functions is still O(K).

Where K is the number of input vectors. Although it is always two in the setting of this problem, this variable becomes relevant once the input becomes K vectors.

Things to watch for

Skipping Empty Vectors in Initialization

When one or both input vectors are empty, failing to handle them causes index errors or incorrect iteration. Empty vectors should either be skipped entirely or handled gracefully in the next() logic.

Incrementing Element Pointer at Wrong Time

In the two-pointer approach, the element index should only increment after completing a full round through all vectors (when p_vec wraps back to 0). Incrementing it every step breaks the zigzag order.

Not Returning After Finding a Valid Element

In the two-pointer approach, after finding a valid element in a vector, you must return immediately. Continuing to iterate through other vectors in the same call corrupts the zigzag ordering and may return wrong elements.

from typing import List
from collections import deque

class ZigzagIterator:
    def __init__(self, v1: List[int], v2: List[int]):
        self.vectors = [v1, v2]
        self.queue = deque()
        for index, vector in enumerate(self.vectors):
            # <index_of_vector, index_of_element_to_output>
            if len(vector) > 0:
                self.queue.append((index, 0))

    def next(self) -> int:

        if self.queue:
            vec_index, elem_index = self.queue.popleft()
            next_elem_index = elem_index + 1
            if next_elem_index < len(self.vectors[vec_index]):
                # append the pointer for the next round
                # if there are some elements left
                self.queue.append((vec_index, next_elem_index))

            return self.vectors[vec_index][elem_index]

        # no more element to output
        raise Exception

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

class ZigzagIterator {
private:
    vector<vector<int>> vectors;
    queue<pair<int, int>> q;

public:
    ZigzagIterator(vector<int>& v1, vector<int>& v2) {
        vectors.push_back(v1);
        vectors.push_back(v2);

        for (int index = 0; index < vectors.size(); index++) {
            if (vectors[index].size() > 0) {
                // <index_to_vec, index_to_element_within_vec>
                q.push({index, 0});
            }
        }
    }

    int next() {
        // <index_to_vec, index_to_element_within_vec>
        pair<int, int> pointer = q.front();
        q.pop();

        int vecIndex = pointer.first;
        int elemIndex = pointer.second;
        int nextElemIndex = elemIndex + 1;

        // append the pointer for the next round
        // if there are some elements left.
        if (nextElemIndex < vectors[vecIndex].size()) {
            q.push({vecIndex, nextElemIndex});
        }

        return vectors[vecIndex][elemIndex];
    }

    bool hasNext() {
        return q.size() > 0;
    }
};