Search code examples
pythoncombinationspython-itertools

Is there a python itertools way of generating this data?


I have the following function:

def e(w):
    for i in range(w):
        for j in range(w-1):
            for k in range(w-2):
                yield [i,j,k]

print([i for i in e(4)]) generates data like this:

[[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [0, 2, 0], [0, 2, 1],
 [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1], [1, 2, 0], [1, 2, 1], 
 [2, 0, 0], [2, 0, 1], [2, 1, 0], [2, 1, 1], [2, 2, 0], [2, 2, 1], 
 [3, 0, 0], [3, 0, 1], [3, 1, 0], [3, 1, 1], [3, 2, 0], [3, 2, 1]]

I have generalised e() into a recursive function, f() :

def f(w, mx=0, depth=0, values=[]):
    if mx==0:
        yield values
    for i in range(w - depth):
        yield from f(w, mx-1, depth+1, values + [i])

print([i for i in f(4,3)]) generates the same data.
Can this sort of data be generated using a standard library like itertools or numpy for instance?


Solution

  • If you want to generalise to arbitrary depth, you can use itertools.product on a sequence of iterables generated by a comprehension:

    import itertools
    
    def combs(w, depth):
        # Returns an iterator object which yields the combinations
        return itertools.product(*(range(w-n) for n in range(depth)))
    
    print(list(combs(4,3)))
    

    which outputs (format tidied up to match yours):

    [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1),
    (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1),
    (2, 0, 0), (2, 0, 1), (2, 1, 0), (2, 1, 1), (2, 2, 0), (2, 2, 1),
    (3, 0, 0), (3, 0, 1), (3, 1, 0), (3, 1, 1), (3, 2, 0), (3, 2, 1)]