5
\$\begingroup\$

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]
\$\endgroup\$
  • 1
    \$\begingroup\$ Not an answer, as it's 'incorrect', but I'd personally use the roundrobin function from the itertools recipes. \$\endgroup\$ – Peilonrayz Mar 30 '17 at 23:56
3
\$\begingroup\$

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)
\$\endgroup\$
2
\$\begingroup\$

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.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.