Skip to content

281. Zigzag Iterator 👍

  • Time:
    • Constructor: $O(1)$
    • next(): $O(1)$
    • hasNext(): $O(1)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class ZigzagIterator {
 public:
  ZigzagIterator(vector<int>& v1, vector<int>& v2) {
    if (!v1.empty())
      q.emplace(v1.begin(), v1.end());
    if (!v2.empty())
      q.emplace(v2.begin(), v2.end());
  }

  int next() {
    const auto [it, endIt] = q.front();
    q.pop();
    if (it + 1 != endIt)
      q.emplace(it + 1, endIt);
    return *it;
  }

  bool hasNext() {
    return !q.empty();
  }

 private:
  // {{ it, endIt }}
  queue<pair<vector<int>::iterator, vector<int>::iterator>> q;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
public class ZigzagIterator {
  public ZigzagIterator(List<Integer> v1, List<Integer> v2) {
    if (!v1.isEmpty())
      q.offer(v1.iterator());
    if (!v2.isEmpty())
      q.offer(v2.iterator());
  }

  public int next() {
    final Iterator it = q.poll();
    final int next = (int) it.next();
    if (it.hasNext())
      q.offer(it);
    return next;
  }

  public boolean hasNext() {
    return !q.isEmpty();
  }

  private Queue<Iterator> q = new ArrayDeque<>();
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class ZigzagIterator:
  def __init__(self, v1: list[int], v2: list[int]):
    def vals():
      for i in itertools.count():
        for v in v1, v2:
          if i < len(v):
            yield v[i]
    self.vals = vals()
    self.n = len(v1) + len(v2)

  def next(self):
    self.n -= 1
    return next(self.vals)

  def hasNext(self):
    return self.n > 0