pythonpython-3.xgeneratorpython-itertools

Python product of infinite generators


I'm trying to get the product of 2 infinite generators but the product function in itertools doesn't allow this sort of behaviour.

Example behaviour:

from itertools import *
i = count(1)
j = count(1)
x = product(i, j)

[Killed]

What I want:

x = product(i, j)

((0,0), (0,1), (1,0), (1,1) ...)

It doesn't matter in which order the combinations get returned as long as given infinite time, all combinations will be eventually generated. This means that given a combination of elements, there must be a finite index in the returned generator with that combination.


Solution

  • tl;dr

    The code presented below is now included in the package infinite on PyPI. So now you can actually just pip install infinite before running this:

    from itertools import count
    from infinite import product
    
    for x, y in product(count(0), count(0)):
        print(x, y)
        if (x, y) == (3, 3):
            break
    

    The zig-zag algorithm

    What is needed is a way to iterate the pairs of numbers so that looking for a specific pair (made of finite numbers) can be done in finite time. A way to go this is the zig-zag scanning algorithm.

    zig-zag scanning algorithm

    In order to do it you need to cache previous values, so I created a class GenCacher to cache previously extracted values:

    class GenCacher:
        def __init__(self, generator):
            self._g = generator
            self._cache = []
    
        def __getitem__(self, idx):
            while len(self._cache) <= idx:
                self._cache.append(next(self._g))
            return self._cache[idx]
    

    After that you just need to implement the zig-zag algorithm:

    def product(gen1, gen2):
        gc1 = GenCacher(gen1)
        gc2 = GenCacher(gen2)
        idx1 = idx2 = 0
        moving_up = True
    
        while True:
            yield (gc1[idx1], gc2[idx2])
    
            if moving_up and idx1 == 0:
                idx2 += 1
                moving_up = False
            elif not moving_up and idx2 == 0:
                idx1 += 1
                moving_up = True
            elif moving_up:
                idx1, idx2 = idx1 - 1, idx2 + 1
            else:
                idx1, idx2 = idx1 + 1, idx2 - 1
    

    Example

    from itertools import count
    
    for x, y in product(count(0), count(0)):
        print(x, y)
        if x == 2 and y == 2:
            break
    

    This produces the following output:

    0 0
    0 1
    1 0
    2 0
    1 1
    0 2
    0 3
    1 2
    2 1
    3 0
    4 0
    3 1
    2 2
    

    Extend the solution to more than 2 generators

    We can edit the solution a bit to make it work even for multiple generators. The basic idea is:

    1. iterate over the distance from (0,0) (the sum of the indexes). (0,0) is the only one with distance 0, (1,0) and (0,1) are at distance 1, etc.

    2. generate all the tuples of indexes for that distance

    3. extract the corresponding element

    We still need the GenCacher class, but the code becomes:

    def summations(sumTo, n=2):
        if n == 1:
            yield (sumTo,)
        else:
            for head in xrange(sumTo + 1):
                for tail in summations(sumTo - head, n - 1):
                    yield (head,) + tail
    
    def product(*gens):
        gens = map(GenCacher, gens)
    
        for dist in count(0):
            for idxs in summations(dist, len(gens)):
                yield tuple(gen[idx] for gen, idx in zip(gens, idxs))