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
hasNextreturns 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
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
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
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]