在Python中生成满足某些约束的所有列表

Question

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?

Answer 1

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!)

Answer 2

You can use itertools.permutations and set (to eliminate duplicates):

>>> from itertools import permutations
>>> for combo in set(permutations([1, 1, 1, 2, 2])):
...     print(list(combo))
...
[1, 2, 1, 1, 2]
[2, 1, 1, 1, 2]
[2, 1, 2, 1, 1]
[2, 1, 1, 2, 1]
[1, 1, 2, 1, 2]
[1, 1, 1, 2, 2]
[1, 2, 1, 2, 1]
[1, 1, 2, 2, 1]
[1, 2, 2, 1, 1]
[2, 2, 1, 1, 1]
>>>

If the combinations need to be in order, then you can use sorted :

>>> for combo in sorted(set(permutations([1, 1, 1, 2, 2]))):
...    print(list(combo))
...
[1, 1, 1, 2, 2]
[1, 1, 2, 1, 2]
[1, 1, 2, 2, 1]
[1, 2, 1, 1, 2]
[1, 2, 1, 2, 1]
[1, 2, 2, 1, 1]
[2, 1, 1, 1, 2]
[2, 1, 1, 2, 1]
[2, 1, 2, 1, 1]
[2, 2, 1, 1, 1]
>>>