Zigzag Iterator的Pythonic方式？

Question

I'm programming onto Zigzag Iterator, it is to iterate a 2D list in following way:

[1,4,7]
[2,5,8,9]
[3,6]

to

[1,2,3,4,5,6,7,8,9]

I implemented an algorithem:

class ZigzagIterator:

    def __init__(self, vecs):

        self.vecs = []
        self.turns = 0
        for vec in vecs:
            vec and self.vecs.append(iter(vec))

    def next(self):
        try:
            elem = self.vecs[self.turns].next()
            self.turns = (self.turns+1) % len(self.vecs)
            return elem
        except StopIteration:
            self.vecs.pop(self.turns)
            if self.hasNext():
                self.turns %= len(self.vecs)

    def hasNext(self):
        return len(self.vecs) > 0

if __name__ == "__main__":
    s = ZigzagIterator([[1,4,7],[2,5,8,9],[3,6]])
    while s.hasNext():
        print s.next()

>>> 1 2 3 4 5 6 7 8 None None 9 None

I know the problem is because I call 1 more time next() of each list, then I get 3 None. I could resolve this issue by checking the hasnext method with java. I can also implement a hasnext Iterator in python. My questions is how I can solve this problem in a more pythonic way rather than thinking it in Java.

Answer 1

This is the round robin recipe in the itertools docs .

def roundrobin(*iterables):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    # Recipe credited to George Sakkis
    num_active = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while num_active:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            # Remove the iterator we just exhausted from the cycle.
            num_active -= 1
            nexts = cycle(islice(nexts, num_active))

Answer 2

This can be easily built with tools in itertools :

from itertools import zip_longest, chain

sentinel = object()

def zigzag(lists):
    return (
        value
        for value
        in chain.from_iterable(zip_longest(*lists, fillvalue=sentinel))
        if value is not sentinel
    )

lists = [
    [1,4,7],
    [2,5,8,9],
    [3,6],
]

print(list(zigzag(lists)))

The sentinel stuff is needed so None values can be safely zigzagged over. (It's a value that's supposed to be guaranteed to not show up in the original lists.)

Answer 3

from itertools import chain, zip_longest
print([i for i in chain.from_iterable(zip_longest([1,4,7], [2,5,8,9], [3,6])) if i is not None])

This outputs:

[1, 2, 3, 4, 5, 6, 7, 8, 9]

Answer 4

For pythonic solution you need to implement the iterator protocol, I hope this is exactly what you were looking for.

from itertools import chain, zip_longest


class ZigZagIterator:
    def __init__(self, *lists):
        self.elements = chain(*zip_longest(*lists))

    def __iter__(self):
        for num in self.elements:
            if num is not None:
                yield num


zig = ZigZagIterator([1, 4, 7], [2, 5, 8, 9], [3, 6])

for num in zig:
    print(num)

If you really want to use has_next and next then

from itertools import chain, zip_longest


class ZigZagIterator:
    def __init__(self, *lists):
        elements = chain(*zip_longest(*lists))
        self.elements = filter(lambda x: x is not None, elements)

    def has_next(self):
        try:
            self.next_value = next(self.elements)
        except StopIteration:
            return False
        return True

    def next(self):
        return self.next_value


zig = ZigZagIterator([1, 4, 7], [2, 5, 8, 9], [3, 6])

while zig.has_next():
    print(zig.next())