在Python中测试所有组合

Question

我有两组选择：

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

这两组选择的数量不一定相同，而且我需要从第一组选择两次。我该如何确保我尝试了第一组中所有可能的两个选项和第二组中的一个选项的组合呢？下面是一个选择的例子...

selectedOptions = [an1,an2,bn]

Answer 1

你可以使用 itertools.product 来实现这个功能。它会返回所有可能的组合。

举个例子：

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

这会生成一些“重复”的组合，比如 [a1,a1,b2] 和 [a2,a3,b2]、[a3,a2,b2]。你可以通过过滤来解决这个问题。下面的代码可以防止出现任何重复的组合*：

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

(*) 这里假设选项有某种自然的顺序，这在所有基本值中都是成立的。

shang 的 回答 也非常不错。我写了一些代码来比较它们：

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

结果是：

0.189330005646
0.428138256073

所以生成器的方法更快。不过，速度并不是一切。就我个人而言，我觉得我的代码更“干净”，但选择权在你！

(顺便说一下，它们的计数是一样的，所以两者都是正确的。)

Answer 2

结合使用 product 和 permutations 这两个来自 itertools 的功能，前提是你不想从第一个列表中得到重复的结果：

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

Answer 3

假设你不想让list1里面有重复的条目，这里有一个生成器可以用来遍历所有的组合：

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

不过，这个方法不会选择list1里面的选项以不同的顺序。如果你想要得到[a1, a2, b1]和[a2, a1, b1]这两种组合，你可以使用：

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