CODEWITHSUNDEEP
pythonalgorithm
Sean Holdsworth
Sean Holdsworth

Reputation: 434

Traversing a sequence of generators

I have a sequence of generators: (gen_0, gen_1, ... gen_n)

These generators will create their values lazily but are finite and will have potentially different lengths.

I need to be able to construct another generator that yields the first element of each generator in order, followed by the second and so forth, skipping values from generators that have been exhausted.

I think this problem is analogous to taking the tuple

((1, 4, 7, 10, 13, 16), (2, 5, 8, 11, 14), (3, 6, 9, 12, 15, 17, 18))

and traversing it so that it would yield the numbers from 1 through 18 in order.

I'm working on solving this simple example using (genA, genB, genC) with genA yielding values from (1, 4, 7, 10, 13, 16), genB yielding (2, 5, 8, 11, 14) and genC yielding (3, 6, 9, 12, 15, 17, 18).

To solve the simpler problem with the tuple of tuples the answer is fairly simple if the elements of the tuple were the same length. If the variable 'a' referred to the tuple, you could use

[i for t in zip(*a) for i in t]

Unfortunately the items are not necessarily the same length and the zip trick doesn't seem to work for generators anyway.

So far my code is horribly ugly and I'm failing to find anything approaching a clean solution. Help?

Upvotes: 5

Views: 165

Answers (4)

Abhijit
Abhijit

Reputation: 63787

I think you need itertools.izip_longest

>>> list([e for e in t if  e is not None] for t in itertools.izip_longest(*some_gen,
                                                               fillvalue=None))
[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15], [16, 17], [18]]
>>>

Upvotes: 8

kaspersky
kaspersky

Reputation: 4117

You might consider itertools.izip_longest, but in case None is a valid value, that solution will fail. Here is a sample "another generator", which does exactly what you asked for, and is pretty clean:

def my_gen(generators):
    while True:
        rez = () 
        for gen in generators:
            try:
                rez = rez + (gen.next(),)
            except StopIteration:
                pass
        if rez:
            yield rez
        else:
            break

print [x for x in my_gen((iter(xrange(2)), iter(xrange(3)), iter(xrange(1))))]

[(0, 0, 0), (1, 1), (2,)] #output

Upvotes: 1

Gareth Rees
Gareth Rees

Reputation: 65894

If you look at the documentation for itertools.izip_longest, you'll see that it gives a pure-Python implementation. It's easy to modify this implementation so that it produces the results you need instead (that is, just like izip_longest, but without any fillvalue):

class ZipExhausted(Exception):
    pass

def izip_longest_nofill(*args):
    """
    Return a generator whose .next() method returns a tuple where the
    i-th element comes from the i-th iterable argument that has not
    yet been exhausted. The .next() method continues until all
    iterables in the argument sequence have been exhausted and then it
    raises StopIteration.

    >>> list(izip_longest_nofill(*[xrange(i,2*i) for i in 2,3,5]))
    [(2, 3, 5), (3, 4, 6), (5, 7), (8,), (9,)]
    """
    iterators = map(iter, args)
    def zip_next():
        i = 0
        while i < len(iterators):
            try:
                yield next(iterators[i])
                i += 1
            except StopIteration:
                del iterators[i]
        if i == 0:
            raise ZipExhausted
    try:
        while iterators:
            yield tuple(zip_next())
    except ZipExhausted:
        pass

This avoids the need to re-filter the output of izip_longest to discard the fillvalues. Alternatively, if you want a "flattened" output:

def iter_round_robin(*args):
    """
    Return a generator whose .next() method cycles round the iterable
    arguments in turn (ignoring ones that have been exhausted). The
    .next() method continues until all iterables in the argument
    sequence have been exhausted and then it raises StopIteration.

    >>> list(iter_round_robin(*[xrange(i) for i in 2,3,5]))
    [0, 0, 0, 1, 1, 1, 2, 2, 3, 4]
    """
    iterators = map(iter, args)
    while iterators:
        i = 0
        while i < len(iterators):
            try:
                yield next(iterators[i])
                i += 1
            except StopIteration:
                del iterators[i]

Upvotes: 4

Joachim Isaksson
Joachim Isaksson

Reputation: 181087

Another itertools option if you want them all collapsed in a single list; this (as @gg.kaspersky already pointed out in another thread) does not handle generated None values.

g = (generator1, generator2, generator3)

res = [e for e in itertools.chain(*itertools.izip_longest(*g)) if e is not None]
print res

[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]

Upvotes: 2

Related Questions