Zigzag Iterator

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):
"""
: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]

• Not an answer, as it's 'incorrect', but I'd personally use the roundrobin function from the itertools recipes. – Peilonrayz Mar 30 '17 at 23:56

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)


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


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