Search code examples
pythonfor-looppython-itertools

Two nested for loops involving itertools don't produce the permutations of the outer loop


I have the following bit of code:

a = [0,1,2,3,4]
b = range(6)
p = itertools.permutations(a)
p2 = itertools.product(b,b,b,b,b)

for pos in p:
    for shift in p2:
        print(pos,shift)

which prints out:

(0, 1, 2, 3, 4) (0, 0, 0, 0, 0)
(0, 1, 2, 3, 4) (0, 0, 0, 0, 1)
(0, 1, 2, 3, 4) (0, 0, 0, 0, 2)
(0, 1, 2, 3, 4) (0, 0, 0, 0, 3)
(0, 1, 2, 3, 4) (0, 0, 0, 0, 4)
...
(0, 1, 2, 3, 4) (5, 5, 5, 5, 5)

The outer loop is executed only once. If I do

for pos in p:
    print(pos)

the loop executes just fine:

(0, 1, 2, 3, 4)
(0, 1, 2, 4, 3)
...
(4, 3, 2, 1, 0)

Why is this happening, and how do I make the two nested for loops work? Note that I have tried looping over list(p) instead of just p and it doesn't make a difference.


Solution

  • The reason for why the code isn't operating as you expect it to is the same reason as why this wouldn't work:

    p = iter(range(6))
    p2 = iter(range(6))
    
    for pos in p:
        for shift in p2:
            print(pos,shift)
    

    itertools.product returns an iterator that can only be consumed once. The first time you loop through p2 you have consumed all of the values, and subsequent iterations of the outer loop essentially become a no-op.

    You need to store the iterator output to a container that can be iterated over multiple times, or call itertools.product multiple times to reconstruct as necessary.

    import itertools
    
    a = [0,1,2,3,4]
    b = range(6)
    p = itertools.permutations(a)
    p2 = list(itertools.product(b,b,b,b,b))
    
    for pos in p:
        for shift in p2:
            print(pos,shift)
    

    Or

    import itertools
    
    a = [0,1,2,3,4]
    b = range(6)
    p = itertools.permutations(a)
    
    for pos in p:
        p2 = itertools.product(b,b,b,b,b)
        for shift in p2:
            print(pos,shift)
    

    Note: In reality, the permutations and product methods are actually constructor calls to generate operation specific classes, however, they do implement the iterator interface so I'm just calling them that!