组合问题以易于陈述但可能难以解决而闻名。 对于这个问题,我不会使用itertools
,而是编写一个自定义生成器。 例如,
def combs(elt2color, combination_size=4, max_colors=3):
def inner(needed, index):
if needed == 0:
yield result
return
if n - index < needed:
return
for _ in inner(needed, index + 1):
yield result
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
if not n or not iterables:
if not n and not iterables:
yield ()
return
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):
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 中进行原型设计的,但在一定程度上,它们被简化为使用小整数向量(不使用集合、字典、序列切片或粘合部分结果序列的“内部循环” - 在初始化后尽可能地坚持对脏简单本机 C 类型的 O(1) 最坏情况时间操作)。
在更高的层面上,这个 itertools
函数将接受任何可迭代对象作为其主要参数,并几乎一定保证按字典索引顺序返回组合。因此,下面是可以做到所有这些的代码。除了 iterable
参数之外,它还需要一个 elt2ec
映射,该映射将可迭代对象中的每个元素映射到其等效类(对于您来说,这些是命名颜色的字符串,但任何可用作字典键的对象都可以用作等效类):
def combs3(iterable, elt2ec, k, maxec):
elts = tuple(iterable)
n = len(elts)
ec = [None] * n
ec2j = {}
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()
中没有比用一个小整数索引向量更高级的东西。唯一剩下的事情是消除递归生成。虽然这很繁琐,但由于它在此处不会说明任何特别有趣的东西,因此我将忽略它。
无论如何,有趣的事情是计时。如评论中所述,计时结果受您使用的测试用例强烈影响。combs3()
在某些情况下最快,但并不经常!它几乎总是比我的原始combs()
快,但通常比我的combs2()
或@GarethRees的可爱的constrained_combinations()
慢。
那么,当combs3()
已经优化到“几乎所有操作都是无意识的C级操作”时,为什么会发生这种情况呢?很容易理解!它仍然是用Python编写的。combs2()
和constrained_combinations()
使用C编码的itertools.combinations()
来完成大部分工作,这使得它们之间存在明显差异。如果combs3()
是用C编写的,那么它将超越它们。
当然,任何一个都可以比原始帖子中的allowed_combinations()
运行得更快 - 但它也可能是最快的(例如,选择几乎所有输入,其中max_colors
非常大以至于没有排除任何组合 - 然后allowed_combinations()
浪费的努力很少,而所有这些其他方法都会增加很多额外的开销来“优化”永远不会发生的修剪)。
constrained_combinations(nodes, 4, 3)
。但是,如果我运行你的测试(即使仍然在笔记本中),那么你是正确的,combs
和constrained_combinations
非常相似。 - alberto