Itertools instead of nested loops

Question

list_a = [] 
for color in [True,False]:
    for piece in range(1,7):
        list_a = list_a + function(piece,color)

Here the function(piece,color) returns a list, which I want to join and finally return the long list, can itertools.chain be used here? because I think it might be faster. I am only displaying an example, but in my actual code the loop runs about 100,000 times, which is why I am looking for a faster method.

Answer 1

I'm going to answer the question you should have asked instead ;-)

This:

list_a = list_a + function(piece,color)

takes time quadratic in the number of times it's executed. Each time, a brand new list object is created, copying the entirety of the old list_a and the new list.

So if it's executed many times, you can get a huge improvement by changing it to this:

list_a.extend(function(piece,color))

Then list_a is extended "in place" whenever possible; under the covers, it may need to make a copy to a larger memory area from time to time, but overall the amortized time is linear in the number of times it's executed.

Answer 2

If you really want to use itertools.chain :

>>> from itertools import product, chain
>>> list_a = list(chain.from_iterable(function(piece, color) for piece, color in product([True, False], range(1, 7))))

Of course, if you use list_a += function(piece, color) this would likely be just as fast (maybe faster?).

The problem with list_a = list_a + function(piece, color) is that this line is quadratic in it's input, because it builds an entirely new list, whereas list_a += function(piece, color) is the equivalent of using extend , which for Python lists is amortized constant time, so the inner part stays linear instead of quadratic.

Answer 3

yield from seems to be the simple solution here.

def generator():
    for color in [True,False]:
        for piece in range(1,7):
            yield from function(piece,color)

Answer 4

You can use itertools.starmap , although you'll still have the loops, you can't really escape from them:

result_generator = starmap(function, ((piece, color) for piece in range(1,7) for color in [True,False]))

list_a = list(result_generator)

Then, use itertools.product like this:

result_generator = starmap(function, product(range(1,7), [True,False]))