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]

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

200_success · Accepted Answer · 2017-03-30 19:05:19Z

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)

Community · Accepted Answer · 2017-05-23 12:41:03Z

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.

Stack Exchange Network

Zigzag Iterator

2 Answers 2

You must log in to answer this question.

Hot Network Questions

Zigzag Iterator

2 Answers 2

You must log in to answer this question.

Related

Hot Network Questions