Question

I want itertools.product to return a list instead of a tuple. I am currently doing it by creating my own function like this:

def product_list(*args, **kwds):
    # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
    # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
    pools = map(tuple, args) * kwds.get('repeat', 1)
    result = [[]]
    for pool in pools:
        result = [x + [y] for x in result for y in pool]
    for prod in result:
        yield list(prod)  # Yields list() instead of tuple()

The code is from the Python docs - I just modified the last line. This works fine but doesn't seem very clever.

What are the other ways of doing this? I'm thinking of using something like a decorator or wrapping it with my own generator function. I'm not too familiar with either concepts so would appreciate if someone can show me.

Edit I'm doing something messy like this:

for r0 in product_list([0, 1], repeat=3):
    r0.insert(0, 0)
    for r1 in product_list([0, 1], repeat=3):
        r1.insert(1, 0)
        for r2 in product_list([0, 1], repeat=3):
            r2.insert(2, 0)
            for r3 in product_list([0, 1], repeat=3):
                r3.insert(3, 0)

so I would prefer my function return a list instead of having to cast it every time. (I know the code is messy and needs recursion but I will think about that later. I'm more interested to learn how to do what I described above)

Was it helpful?

Solution

itertools.product is a generator, and you can easily chain generators together. Here's a generator expression that changes each tuple yielded by product into a list:

(list(tup) for tup in itertools.product(iterable1, iterable2, etc))

In your example code, you could use the generator expression, or you could use a different approach to add an extra value on to the front of your values while keeping them as tuples:

for r0 in itertools.product([0, 1], repeat=3):
    r0 = (0,) + r0 # keep r0 a tuple!
    for r1 in itertools.product([0, 1], repeat=3):
        r1 = (1,) + r1 # same here
        # ...

Since you don't show what you're using your rN variables for, it's impossible to give you a definitive answer as to what would be the best way to go. (It is a bit of a code smell that you have numbered variables.) Indeed, since your loops are simply computing three more 0 or 1 digits, you might be able to get away with a single product call, that generates an list of n different r values in one go:

for bits in itertools.product([0, 1], repeat=3*n):
    rs = [(i,) + bits[3*i:3*i+3] for i in range(n)]
    # do something with the list of r tuples here
Licensed under: CC-BY-SA with attribution
Not affiliated with StackOverflow
scroll top