Skip to content

281. Zigzag Iterator

On LeetCode ->

Problem

Implement an iterator over two integer lists: alternate elements starting with the first list, skipping exhausted lists. Provide next() and non-consuming hasNext().

For \(k\) lists, visit nonempty lists cyclically.

v1 = [1, 2], v2 = [3, 4, 5]
Repeated next(): 1 → 3 → 2 → 4 → 5
hasNext(): False

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

    def hasNext(self):
        pass

    def next(self):
        pass

Key trick

Keep a queue of nonempty lists and their positions; after consuming an element, move that list to the back only if it has elements remaining.

Trap

Do not stop when one list ends, enqueue empty lists, or let repeated hasNext() calls consume elements.

Why is it interesting?

It combines iterator state with round-robin scheduling, and the queue approach extends directly to \(k\) lists.

Python solution

from collections import deque


class ZigzagIterator:
    def __init__(self, v1: list[int], v2: list[int]):
        # Only lists with remaining elements participate in the rotation.
        self.q = deque((vec, 0) for vec in (v1, v2) if vec)

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

    def next(self) -> int:
        if not self.q:
            raise StopIteration

        vec, i = self.q.popleft()
        val = vec[i]
        i += 1

        if i < len(vec):
            self.q.append((vec, i))

        return val

Complexity: Initialization, hasNext(), and next() take \(O(1)\) time; auxiliary space is \(O(1)\) for two lists. For \(k\) lists, initialization takes \(O(k)\) time and auxiliary space is \(O(k)\); each operation remains \(O(1)\).

Comment on my solution

Your solution is correct, with \(O(1)\) time per operation and \(O(1)\) auxiliary space. Although hasNext() changes cur_v, repeated checks do not affect the output; the queue approach makes hasNext() read-only and avoids scanning exhausted lists when generalized to \(k\) vectors.

class ZigzagIterator:
    def __init__(self, v1: list[int], v2: list[int]):
        self.vec = [v1, v2]
        self.cur_v = 0
        self.pos = [0, 0]

    def hasNext(self):
        if self.pos[self.cur_v] < len(self.vec[self.cur_v]):
            return True
        self.cur_v = (self.cur_v + 1) % 2
        if self.pos[self.cur_v] < len(self.vec[self.cur_v]):
            return True
        return False

    def next(self):
        if not self.hasNext():
            raise StopIteration
        v = self.cur_v
        pos = self.pos[v]
        val = self.vec[v][pos]
        self.pos[v] += 1
        self.cur_v = (self.cur_v + 1) % 2
        return val

Extra

Problem description

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.

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]

Constraints:

  • 0 <= v1.length, v2.length <= 1000
  • 1 <= v1.length + v2.length <= 2000
  • -2^31 <= v1[i], v2[i] <= 2^31 - 1

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

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".

Follow-up Example:

  Input: v1 = [1,2,3], v2 = [4,5,6,7], v3 = [8,9]
  Output: [1,4,8,2,5,9,3,6,7]

What is round-robin scheduling?

What is round-robin scheduling?

Round-robin scheduling means giving each active participant one turn, in a fixed repeating order. After the last participant, return to the first.

For your iterator, each list is a participant, and a turn means returning one element:

A = [1, 2, 3]
B = [4]
C = [5, 6]

Turns:   A → B → C → A → C → A
Output:  1 → 4 → 5 → 2 → 6 → 3

Once a list is exhausted, it stops participating.

A queue naturally represents whose turn comes next:

Queue before turn    Action                      Queue afterward
[A, B, C]            Take from A; A has more      [B, C, A]
[B, C, A]            Take from B; B is exhausted  [C, A]
[C, A]               Take from C; C has more      [A, C]

The important distinction is one item per turn, not consuming an entire list before moving on. This gives every active list regular opportunities to contribute.

In operating systems, the same idea schedules programs: each runnable program gets a short CPU time slice, then moves to the back of the queue if it still needs to run.