Testing all combinations in Python

问题:

I have two sets of options:

optionList1 = [a1,a2,a3,...,an]
optionList2 = [b1,b2,b3,...,bn]

The number of elements in the optionlists are not necessarily equal and I have to choose from the first optionlist twice. How do I ensure that I have tried every combination of 2 options from the first list and one from the second list. An example selection set below...

selectedOptions = [an1,an2,bn]

回答1:

Assuming you don't want duplicate entries from list1, here's a generator that you can use to iterate over all combinations:

def combinations(list1, list2):
    return ([opt1, opt2, opt3]
            for i,opt1 in enumerate(list1)
            for opt2 in list1[i+1:]
            for opt3 in list2)

This, however, doesn't select the same options from list1 in different orderings. If you want to get both [a1, a2, b1] and [a2, a1, b1] then you can use:

def combinations(list1, list2):
    return ([opt1, opt2, opt3]
            for opt1 in list1
            for opt2 in list1
            for opt3 in list2 if opt1 != opt2)

回答2:

Combine product and permutations from itertools assuming you don't want duplicates from the first list:

>>> from itertools import product,permutations
>>> o1 = 'a1 a2 a3'.split()
>>> o2 = 'b1 b2 b3'.split()
>>> for (a,b),c in product(permutations(o1,2),o2):
...     print a,b,c
... 
a1 a2 b1
a1 a2 b2
a1 a2 b3
a1 a3 b1
a1 a3 b2
a1 a3 b3
a2 a1 b1
a2 a1 b2
a2 a1 b3
a2 a3 b1
a2 a3 b2
a2 a3 b3
a3 a1 b1
a3 a1 b2
a3 a1 b3
a3 a2 b1
a3 a2 b2
a3 a2 b3

回答3:

You can use itertools.product for this. It returns all possible combinations.

For example

for a1, a2, b in itertools.product(optionlist1,optionlist1,optionlist2):
    do_something(a1,a2,b)

This will produce "doubles" as [a1,a1,b2] and [a2,a3,b2],[a3,a2,b2]. You can fix this with a filter. The following prevents any doubles*:

for a1,a2,b in itertools.ifilter(lambda x: x[0]<x[1], itertools.product(optionlist1,optionlist1,optionlist2)):
    do_something(a1,a2,b)

(*) This assumes that the options have some natural ordering which will be the case with all primitive values.

shang's answer is also very good. I wrote some code to compare them:

from itertools import ifilter, product
import random
from timeit import repeat

def generator_way(list1, list2):
    def combinations(list1, list2):
        return ([opt1, opt2, opt3]
                for i,opt1 in enumerate(list1)
                for opt2 in list1[i+1:]
                for opt3 in list2)
    count = 0
    for a1,a2,b in combinations(list1,list2):
        count += 1

    return count

def itertools_way(list1,list2):
    count = 0
    for a1,a2,b in ifilter(lambda x: x[0] < x[1], product(list1,list1,list2)):
        count += 1
    return count

list1 = range(0,100)
random.shuffle(list1)
list2 = range(0,100)
random.shuffle(list2)

print sum(repeat(lambda: generator_way(list1,list2),repeat = 10, number=1))/10
print sum(repeat(lambda: itertools_way(list1,list2),repeat = 10, number=1))/10

And the result is:

0.189330005646
0.428138256073

So the generator method is faster. However, speed is not everything. Personally I find my code 'cleaner', but the choice is yours!

(Btw, they give both identical counts, so both are equally correct.)

回答4:

It sounds to me like you're looking for itertools.product()

>>> options_a = [1,2]
>>> options_b = ['a','b','c']
>>> list(itertools.product(options_a, options_a, options_b))
[(1, 1, 'a'),
 (1, 1, 'b'),
 (1, 1, 'c'),
 (1, 2, 'a'),
 (1, 2, 'b'),
 (1, 2, 'c'),
 (2, 1, 'a'),
 (2, 1, 'b'),
 (2, 1, 'c'),
 (2, 2, 'a'),
 (2, 2, 'b'),
 (2, 2, 'c')]

回答5:

One way to do is to use itertools.product.

for x, y, z in itertools.product(optionlist1,optionlist1,optionlist2):
    print x,y,z