Problem

Given two vectors of integers v1 and v2, implement an iterator to return their elements alternately.

Implement the ZigzagIterator class:

  • ZigzagIterator(List<int> v1, List<int> v2) initializes the object with the two vectors v1 and v2.
  • boolean hasNext() returns true if the iterator still has elements, and false otherwise.
  • int next() returns the current element of the iterator and moves the iterator to the next element.

Follow up: What if you are given k 1d vectors? How well can your code be extended to such cases?

Examples

Example 1:

Input: v1 = [1,2], v2 = [3,4,5,6]
Output: [1,3,2,4,5,6]
Explanation: By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1,3,2,4,5,6].

Example 2:

Input: v1 = [1], v2 = []
Output: [1]

Example 3:

Input: v1 = [], v2 = [1]
Output: [1]

Clarification for the follow up question The “Zigzag” order is not clearly defined and is ambiguous for k > 2 cases. If “Zigzag” does not look right to you, replace “Zigzag” with “Cyclic”. For example:

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

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

Solution

Method 1 - Set iterators on both vectors

To implement the ZigzagIterator class:

  1. Initialization: Initialize the constructor with the two vectors and set up necessary pointers to track the current position in both vectors.
  2. hasNext Method: This method checks if there are any remaining elements in either of the vectors.
  3. next Method: This method returns the current element and advances the pointer to the next element. The iteration should alternate between the two vectors.

Code

Java
public class ZigzagIterator {
    private List<Integer> v1, v2;
    private int i, j;
    private boolean turn;

    public ZigzagIterator(List<Integer> v1, List<Integer> v2) {
        this.v1 = v1;
        this.v2 = v2;
        this.i = 0;
        this.j = 0;
        this.turn = true; // Start with vector v1
    }

    public boolean hasNext() {
        return i < v1.size() || j < v2.size();
    }

    public int next() {
        if (!hasNext()) {
            throw new RuntimeException("No more elements");
        }

        int nextValue;
        if (turn) {
            if (i < v1.size()) {
                nextValue = v1.get(i++);
            } else {
                nextValue = v2.get(j++);
            }
        } else {
            if (j < v2.size()) {
                nextValue = v2.get(j++);
            } else {
                nextValue = v1.get(i++);
            }
        }

        turn = !turn;
        return nextValue;
    }
}
Python
class ZigzagIterator:
    def __init__(self, v1: List[int], v2: List[int]):
        self.v1 = v1
        self.v2 = v2
        self.i = 0
        self.j = 0
        self.turn = True  # Start with vector v1

    def hasNext(self) -> bool:
        return self.i < len(self.v1) or self.j < len(self.v2)

    def next(self) -> int:
        if not self.hasNext():
            raise Exception("No more elements")

        if self.turn:
            if self.i < len(self.v1):
                next_value = self.v1[self.i]
                self.i += 1
            else:
                next_value = self.v2[self.j]
                self.j += 1
        else:
            if self.j < len(self.v2):
                next_value = self.v2[self.j]
                self.j += 1
            else:
                next_value = self.v1[self.i]
                self.i += 1

        self.turn = not self.turn
        return next_value

Complexity

  • ⏰ Time complexity: O(1) for hasNext and next operations on average.
  • 🧺 Space complexity: O(1) for the iterator itself.