组合问题以易于陈述但可能难以解决而臭名昭著。对于这个,我根本不会使用itertools,而是编写一个自定义生成器。例如,
def combs(elt2color, combination_size=4, max_colors=3):
def inner(needed, index):
if needed == 0:
yield result
return
if n - index < needed:
# not enough elements remain to reach
# combination_size
return
# first all results that don't contain elts[index]
for _ in inner(needed, index + 1):
yield result
# and then all results that do contain elts[index]
needed -= 1
elt = elts[index]
color = elt2color[elt]
color_added = color not in colors_seen
colors_seen.add(color)
if len(colors_seen) <= max_colors:
result[needed] = elt
for _ in inner(needed, index + 1):
yield result
if color_added:
colors_seen.remove(color)
elts = tuple(elt2color)
n = len(elts)
colors_seen = set()
result = [None] * combination_size
for _ in inner(combination_size, 0):
yield tuple(result)
然后:
elt2color = dict([('A', 'red'), ('B', 'red'), ('C', 'blue'),
('D', 'blue'), ('E', 'green'), ('F', 'green'),
('G', 'green'), ('H', 'yellow'), ('I', 'white'),
('J', 'white'), ('K', 'black')])
for c in combs(elt2color):
for element in c:
print("%s-%s" % (element, elements[element]))
print "\n"
产生与您的后处理代码相同的 188 种组合,但在内部放弃部分组合时,它会跨越超过 max_colors 颜色。无法更改 itertools 函数在内部执行的操作,因此当您想要控制它时,您需要自己动手。
使用 itertools
这是另一种方法,首先生成恰好 1 种颜色的所有解决方案,然后恰好生成 2 种颜色,依此类推。 itertools 可以直接用于其中大部分,但在最低级别仍需要自定义生成器。我发现这比完全自定义的生成器更难理解,但对你来说可能更清楚:
def combs2(elt2color, combination_size=4, max_colors=3):
from collections import defaultdict
from itertools import combinations
color2elts = defaultdict(list)
for elt, color in elt2color.items():
color2elts[color].append(elt)
def at_least_one_from_each(iterables, n):
if n < len(iterables):
return # impossible
if not n or not iterables:
if not n and not iterables:
yield ()
return
# Must have n - num_from_first >= len(iterables) - 1,
# so num_from_first <= n - len(iterables) + 1
for num_from_first in range(1, min(len(iterables[0]) + 1,
n - len(iterables) + 2)):
for from_first in combinations(iterables[0],
num_from_first):
for rest in at_least_one_from_each(iterables[1:],
n - num_from_first):
yield from_first + rest
for numcolors in range(1, max_colors + 1):
for colors in combinations(color2elts, numcolors):
# Now this gets tricky. We need to pick
# combination_size elements across all the colors, but
# must pick at least one from each color.
for elements in at_least_one_from_each(
[color2elts[color] for color in colors],
combination_size):
yield elements
我没有对这些进行计时,因为我不在乎 ;-) 完全自定义生成器的单个 result 列表被重复用于构建每个输出,从而降低了动态内存周转率。第二种方法通过将多个级别的from_first 和rest 元组粘贴在一起会产生大量内存流失——这几乎是不可避免的,因为它使用itertools 在每个级别生成from_first 元组。
在内部,itertools 函数几乎总是以更类似于第一个代码示例的方式工作,并且出于同样的原因,尽可能重用内部缓冲区。
还有一个
这更多是为了说明一些微妙之处。我想如果我要在 C 中将这个功能实现为 itertools 函数,我会怎么做。所有itertools 函数首先在 Python 中进行原型设计,但以半低级方式,简化为使用小整数向量(没有“内循环”使用集合、字典、序列切片或将部分结果粘贴在一起)序列 - 尽可能坚持O(1) 在初始化后对简单的原生 C 类型进行最坏情况时间操作。
在更高的级别上,itertools 函数将接受任何可迭代作为其主要参数,并且几乎可以肯定地保证从字典索引顺序返回组合。所以这里的代码可以完成所有这些。除了iterable 参数之外,它还需要一个elt2ec 映射,它将每个元素从可迭代对象映射到其等价类(对您来说,这些是命名颜色的字符串,但任何可用作字典键的对象都可以 用作等价类):
def combs3(iterable, elt2ec, k, maxec):
# Generate all k-combinations from `iterable` spanning no
# more than `maxec` equivalence classes.
elts = tuple(iterable)
n = len(elts)
ec = [None] * n # ec[i] is equiv class ordinal of elts[i]
ec2j = {} # map equiv class to its ordinal
for i, elt in enumerate(elts):
thisec = elt2ec[elt]
j = ec2j.get(thisec)
if j is None:
j = len(ec2j)
ec2j[thisec] = j
ec[i] = j
countec = [0] * len(ec2j)
del ec2j
def inner(i, j, totalec):
if i == k:
yield result
return
for j in range(j, jbound[i]):
thisec = ec[j]
thiscount = countec[thisec]
newtotalec = totalec + (thiscount == 0)
if newtotalec <= maxec:
countec[thisec] = thiscount + 1
result[i] = j
yield from inner(i+1, j+1, newtotalec)
countec[thisec] = thiscount
jbound = list(range(n-k+1, n+1))
result = [None] * k
for _ in inner(0, 0, 0):
yield (elts[i] for i in result)
(注意这是 Python 3 代码。)正如所宣传的那样,inner() 中没有什么比用一个小整数索引一个向量更有趣的了。让它直接翻译成 C 的唯一剩下的就是删除递归生成。这很乏味,因为它不会在这里说明任何特别有趣的事情,所以我将忽略它。
不管怎样,有趣的是计时。如评论中所述,计时结果受您使用的测试用例的强烈影响。 combs3() 这里有时最快,但不经常!它几乎总是比我原来的combs() 快,但通常比我的combs2() 或@GarethRees 的可爱constrained_combinations() 慢。
那么当combs3() 已经被优化“几乎一直到无意识的;-) C 级操作”时,怎么会这样呢?简单!它仍然是用 Python 编写的。 combs2() 和 constrained_combinations() 使用 C 编码的 itertools.combinations() 完成大部分工作,这让世界变得与众不同。 combs3() 会绕着他们转圈如果它是用 C 编码的。
当然,这些中的任何一个都可以比原始帖子中的allowed_combinations() 运行得更快——但那个也可以是最快的(例如,选择max_colors 如此大以至于不排除任何组合的任何输入- 然后allowed_combinations() 几乎不会浪费任何精力,而所有这些其他方法都会增加额外的大量额外开销来“优化”从未发生过的修剪)。