Zigzag Iterator

Question

I've recently solved the LeetCode's "ZigZag iterator" problem:

Given two 1d vectors, implement an iterator to return their elements alternately.

For example, given two 1d vectors:

v1 = [1, 2] 
v2 = [3, 4, 5, 6]

By calling next repeatedly until hasNext returns false, the order of elements returned by next should be:

[1, 3, 2, 4, 5, 6]

Note that there is a pre-defined ZigzagIterator class definition contract that requires to implement __init__(), next() and hasNext() methods.

---

The idea behind the below implementation was to make it work for any number of potential input iterators that need to be iterated in a "zigzag" fashion. That's why I used the izip_longest() and chain():

from itertools import izip_longest, chain
class ZigzagIterator(object):
 def __init__(self, *iterators):
 """
 Initialize your data structure here.
 :type v1: List[int]
 :type v2: List[int]
 """
 self.iterators = chain(*izip_longest(*iterators))
 def next(self):
 """
 :rtype: int
 """
 result = next(self.iterators)
 return self.next() if result is None else result
 def hasNext(self):
 """
 :rtype: bool
 """
 try:
 peek = self.next()
 self.iterators = chain(iter([peek]), self.iterators)
 return True
 except StopIteration:
 return False

This works and passes the LeetCode's OJ tests, but I'm not quite happy with the solution regarding handling the None values created by izip_longest() and peeking into the next value by advancing the iterator and creating a new one, "pre-pending" the peeked value.

What would you improve in the presented code? Is there a better, more optimal approach?

---

FYI, here is a sample ZigzagIterator usage:

v1 = iter([1, 2])
v2 = iter([3, 4, 5, 6])
i, v = ZigzagIterator(v1, v2), []
while i.hasNext():
 v.append(i.next())
print(v) # prints [1, 3, 2, 4, 5, 6]

• Peilonrayz: Not an answer, as it's 'incorrect', but I'd personally use the roundrobin function from the itertools recipes.

Answer 1

This recursion in your next() method is bizarre:

return self.next() if result is None ...

I would just let it fail with a StopIteration.

Furthermore, I would avoid having hasNext() call self.next(), as the whole point of hasNext() is to determine whether it is "safe" to call self.next().

I don't think you need to call iter(...) explicitly within hasNext().

With those remarks, and a few minor simplifications:

from itertools import izip_longest, chain
class ZigzagIterator(object):
 def __init__(self, *iterators):
 self.iter = chain(*izip_longest(*iterators))
 def hasNext(self):
 try:
 self.iter = chain([next(self.iter)], self.iter)
 return True
 except StopIteration:
 return False
 def next(self):
 return next(self.iter)

Answer 2

The problem on LeetCode says you should implement hasNext and next but this is not usually how it is done in Python. I would consider that you revise your code to implement next (and when you decide to use Python 3, __next__) and __iter__. Typically the method hasNext is merged with next. You can see a few examples here:

class Counter:
 def __init__(self, low, high):
 self.current = low
 self.high = high
 def __iter__(self):
 return self
 def next(self): # Python 3: def __next__(self)
 if self.current > self.high:
 raise StopIteration
 else:
 self.current += 1
 return self.current - 1
for c in Counter(3, 8):
 print c

(copied from the first answer)

For the learning experience, instead of giving you the code, I'll leave it for you to re-implement.