HiveBrain v1.2.0
Get Started
← Back to all entries
patternpythonMinor

Zigzag Iterator

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
zigzagiteratorstackoverflow

Problem

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,

Solution

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)

Code Snippets

return self.next() if result is None …
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)

Context

StackExchange Code Review Q#159376, answer score: 3

Revisions (0)

No revisions yet.