Search code examples
pythonlistpython-itertools

Generating all lists that satisfy certain constraints in Python


I would like to generate the following lists in Python:

[1, 1, 1, 2, 2]
[1, 1, 2, 1, 2]

... etc

[2, 1, 2, 1, 1]
[2, 2, 1, 1, 1]

There are always two "2"s and three "1"s in any list.

My intuition suggests that I will need to use the itertools module to do this. However, I am not sure where to begin, though I have read the documentation and looked at examples. Any suggestions?


Solution

  • You can notice that the number of such lists is equal to the number of ways to place two "2"s in a sequence of length 5. This suggests the following solution:

    n = 5 # total length
    n2 = 2 # number of "2"s
    for idx in itertools.combinations( xrange(n), n2 ):
        print [ 2 if i in idx else 1 for i in xrange(n) ]
    

    It's easy to see that the answer using permutations is iterating over n! solutions, while my solution iterates over n!/( (n-n2)! * n2!). For example if the input list is [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2], the solution using permutations is ~90,000,000 times slower (10! * 4!)