Before attempting this problem, you should be comfortable with:
next() and hasNext() methods to traverse elements sequentiallyTo iterate through multiple vectors in zigzag order, we need to alternate between vectors while tracking our position in each. The key insight is to maintain two pointers: one for the current vector and one for the current element index. We cycle through vectors in round-robin fashion, and when we complete a full round across all vectors, we advance the element index. This ensures we visit elements in the correct zigzag order while handling vectors of different lengths.
p_vec for the current vector index and p_elem for the current element index.hasNext() method.next():p_elem, capture it as the return value.p_vec to the next vector (wrapping around using modulo).p_vec wraps back to 0, increment p_elem to move to the next round.hasNext(), return true if the output count is less than the total number of elements.class ZigzagIterator:
def __init__(self, v1: List[int], v2: List[int]):
self.vectors = [v1, v2]
self.p_elem = 0 # pointer to the index of element
self.p_vec = 0 # pointer to the vector
# variables for hasNext() function
self.total_num = len(v1) + len(v2)
self.output_count = 0
def next(self) -> int:
iter_num = 0
ret = None
# Iterate over the vectors
while iter_num < len(self.vectors):
curr_vec = self.vectors[self.p_vec]
if self.p_elem < len(curr_vec):
ret = curr_vec[self.p_elem]
iter_num += 1
self.p_vec = (self.p_vec + 1) % len(self.vectors)
# increment the element pointer once iterating all vectors
if self.p_vec == 0:
self.p_elem += 1
if ret is not None:
self.output_count += 1
return ret
# no more element to output
raise Exception
def hasNext(self) -> bool:
return self.output_count < self.total_numTime Complexity:
For the next() function, at most it will take us iterations to find a valid element to output. Hence, its time complexity is .
For the hasNext() function, its time complexity is .
Space Complexity:
For the next() function, we keep the references to all the input vectors in the variable self.vectors.
As a result, we would need space for vectors.
In addition, we used some constant-space variables such as the pointers to the vector and the element.
Hence, the overall space complexity for this function is .
Note: we did not copy the input vectors, but simply keep references to them.
Where 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 vectors.
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.
index, element index) pairs.next():indices.hasNext(), simply check if the queue is non-empty.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) > 0Time Complexity:
next() function and the hasNext() function, we have a constant time complexity, as we discussed before.Space Complexity:
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 space for 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 .
Where 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 vectors.
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.
# Wrong: Adding empty vectors to queue causes issues
for index, vector in enumerate(self.vectors):
self.queue.append((index, 0)) # Should check: if len(vector) > 0In 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.
# Wrong: Incrementing p_elem on every iteration
self.p_vec = (self.p_vec + 1) % len(self.vectors)
self.p_elem += 1 # Should only increment when p_vec == 0In 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.
# Wrong: Missing early return
if self.p_elem < len(curr_vec):
ret = curr_vec[self.p_elem]
# Missing: return ret after incrementing output_count!
self.p_vec = (self.p_vec + 1) % len(self.vectors)