Want Help Cracking FAANG?

(Then click this)

×
Back to Question Bank

281. Zigzag Iterator - Leetcode Solution

Code Implementation

class ZigzagIterator:
    def __init__(self, v1, v2):
        self.queue = []
        if v1:
            self.queue.append((v1, 0))
        if v2:
            self.queue.append((v2, 0))

    def next(self):
        if not self.hasNext():
            raise Exception("No more elements")
        vec, idx = self.queue.pop(0)
        val = vec[idx]
        if idx + 1 < len(vec):
            self.queue.append((vec, idx + 1))
        return val

    def hasNext(self):
        return len(self.queue) > 0
      
class ZigzagIterator {
public:
    queue<pair<vector<int>*, int>> q;
    ZigzagIterator(vector<int>& v1, vector<int>& v2) {
        if (!v1.empty()) q.push({&v1, 0});
        if (!v2.empty()) q.push({&v2, 0});
    }

    int next() {
        auto p = q.front(); q.pop();
        vector<int>* vec = p.first;
        int idx = p.second;
        int val = (*vec)[idx];
        if (idx + 1 < vec->size()) {
            q.push({vec, idx + 1});
        }
        return val;
    }

    bool hasNext() {
        return !q.empty();
    }
};
      
import java.util.*;

public class ZigzagIterator {
    Queue<Iterator<Integer>> queue;

    public ZigzagIterator(List<Integer> v1, List<Integer> v2) {
        queue = new LinkedList<>();
        if (!v1.isEmpty()) queue.offer(v1.iterator());
        if (!v2.isEmpty()) queue.offer(v2.iterator());
    }

    public int next() {
        Iterator<Integer> it = queue.poll();
        int val = it.next();
        if (it.hasNext()) queue.offer(it);
        return val;
    }

    public boolean hasNext() {
        return !queue.isEmpty();
    }
}
      
var ZigzagIterator = function(v1, v2) {
    this.queue = [];
    if (v1.length > 0) this.queue.push([v1, 0]);
    if (v2.length > 0) this.queue.push([v2, 0]);
};

ZigzagIterator.prototype.next = function() {
    if (!this.hasNext()) throw new Error("No more elements");
    const [vec, idx] = this.queue.shift();
    const val = vec[idx];
    if (idx + 1 < vec.length) this.queue.push([vec, idx + 1]);
    return val;
};

ZigzagIterator.prototype.hasNext = function() {
    return this.queue.length > 0;
};
      

Problem Description

The Zigzag Iterator problem asks you to design an iterator that alternately returns elements from two input lists, v1 and v2. On each call to next(), the iterator should return the next element from the two lists in zigzag (alternating) order. If one list is exhausted before the other, continue iterating through the remaining list. Each element should be returned exactly once, and the iterator should provide a hasNext() method to check if there are more elements to return.

Key constraints:

  • Elements must be returned in alternating order as much as possible.
  • Once an element is returned, it cannot be reused.
  • If one list runs out, continue with the other.
  • There should be no duplicate or skipped elements.

Thought Process

At first glance, the problem may seem like a simple iteration over two lists. However, the alternating requirement means you cannot simply concatenate or merge the lists. A brute-force approach might involve using two pointers and switching between them, but this quickly gets complicated if one list is longer than the other.

To manage the alternation, we need a way to remember which list to pick from next, and also handle the case when one list is exhausted. This leads to the idea of using a queue or similar structure to keep track of which lists have elements left. Each time we return an element, we check if there are more elements in that list, and if so, put it back in the queue for future turns.

This approach generalizes well: not only does it work for two lists, but it can easily be extended to more lists if needed. This insight helps us avoid complicated conditional logic and makes the implementation clean and efficient.

Solution Approach

Let's break down the solution step by step:

  1. Initialization:
    • We use a queue to keep track of which lists have elements left to return.
    • For each non-empty input list, we add a tuple to the queue containing the list itself and the current index (starting at 0).
  2. next() Method:
    • Pop the first element from the queue. This gives us the next list and index to use.
    • Return the element at the current index from that list.
    • If the list has more elements, push it back into the queue with the index incremented by one.
  3. hasNext() Method:
    • Simply check if the queue is not empty. If it is, there are no more elements to return.

This design is efficient because each element is processed exactly once, and the queue ensures we always know which list to draw from next in the zigzag order.

Example Walkthrough

Let's consider the input lists: v1 = [1, 2] and v2 = [3, 4, 5, 6].

  1. Initialization:
    • Queue starts as [([1, 2], 0), ([3, 4, 5, 6], 0)]
  2. First call to next():
    • Pop ([1, 2], 0), return 1. Since index 1 is valid, push ([1, 2], 1) back.
    • Queue: [([3, 4, 5, 6], 0), ([1, 2], 1)]
  3. Second call to next():
    • Pop ([3, 4, 5, 6], 0), return 3. Since index 1 is valid, push ([3, 4, 5, 6], 1) back.
    • Queue: [([1, 2], 1), ([3, 4, 5, 6], 1)]
  4. Third call to next():
    • Pop ([1, 2], 1), return 2. Index 2 is out of bounds, do not push back.
    • Queue: [([3, 4, 5, 6], 1)]
  5. Subsequent calls:
    • Continue popping from queue and returning elements 4, 5, 6 from v2, until the queue is empty.
  6. Final output:
    • Sequence returned: 1, 3, 2, 4, 5, 6

This matches the zigzag order: alternate as long as possible, then finish the remaining list.

Time and Space Complexity

Brute-force approach:

  • Manually switch pointers and check bounds every time, potentially O(n) time per operation due to repeated checks.
  • Space: O(1) extra, but logic is complex and inefficient.
Optimized queue-based approach (our solution):
  • Each next() and hasNext() operation is O(1) time, since we only pop/push from a queue and move pointers.
  • Total time is O(m + n), where m and n are the lengths of the input lists, since each element is processed exactly once.
  • Space complexity is O(k), where k is the number of input lists (here, 2) for the queue, plus O(1) for pointers.

Summary

The Zigzag Iterator problem demonstrates the value of using a queue to manage alternating iteration over multiple lists. By storing the current position of each list in the queue and processing them in turn, we achieve a clean, efficient, and easily extensible solution. The approach ensures each element is returned exactly once in the correct order, and the logic is simple enough for easy maintenance and extension to more lists if needed.