python - 有颜色数量限制的 N 种有色元素的有效组合

Question

给定一组用 C 种颜色着色的 N 个元素，我如何找到长度为 L 且最多包含 M 种颜色的所有可能组合？

我尝试了这个算法，它使用 itertools.combinations 来生成所有可能的组合，然后过滤掉那些不符合最大颜色条件的组合。

from itertools import combinations as cb

def allowed_combinations(elements, combination_size=4, max_colors=3):

    colors = set([c for k, c in elements.items()])
    combinations = cb(elements, combination_size)
    for combination in combinations:
        colors = set([elements[element] for element in combination])
        if len(colors) > max_colors:
            continue
        yield combination


elements = dict()
elements['A'] = 'red'
elements['B'] = 'red'
elements['C'] = 'blue'
elements['D'] = 'blue'
elements['E'] = 'green'
elements['F'] = 'green'
elements['G'] = 'green'
elements['H'] = 'yellow'
elements['I'] = 'white'
elements['J'] = 'white'
elements['K'] = 'black'

combinations = allowed_combinations(elements)

for c in combinations:
    for element in c:
        print("%s-%s" % (element, elements[element]))
    print "\n"

输出如下:

A-red
C-blue
B-red
E-green


A-red
C-blue
B-red
D-blue


A-red
C-blue
B-red
G-green


A-red
C-blue
B-red
F-green

...

问题是生成所有可能的组合在计算上可能非常昂贵。例如，在我的例子中，L 通常为 6，元素数量 N 约为 50，因此它为我们提供了 Bin(50,6) = 15890700 种可能的组合。如果组合中允许的最大颜色数很小，则大多数组合都是“无用的”，因此它们会在过滤步骤中被丢弃。我的直觉是我应该将过滤步骤放在组合步骤之内/之前，以避免组合爆炸，但我不知道该怎么做。

Answer 1

这是一个比目前发布的其他答案更简单的实现。基本方法是:

1. 选择一个目前尚未被选择的值(在您的术语中为“颜色”)；
2. 遍历 i，与将包含在输出中的值关联的键(“元素”)的数量；
3. 遍历 c，将那些长度为 i 的键组合起来；
4. 递归选择下一个值。

from collections import defaultdict, deque
from itertools import combinations

def constrained_combinations(elements, r, s):
    """Generate distinct combinations of 'r' keys from the dictionary
    'elements' using at most 's' different values. The values must be
    hashable.

        >>> from collections import OrderedDict
        >>> elements = OrderedDict(enumerate('aabbc'))
        >>> cc = constrained_combinations
        >>> list(cc(elements, 2, 1))
        [(0, 1), (2, 3)]
        >>> list(cc(elements, 3, 2))
        [(0, 1, 2), (0, 1, 3), (0, 1, 4), (0, 2, 3), (1, 2, 3), (2, 3, 4)]
        >>> list(cc(elements, 3, 3)) == list(combinations(range(5), 3))
        True
        >>> sum(1 for _ in cc(OrderedDict(enumerate('aabbcccdeef')), 4, 3))
        188

    """
    # 'value_keys' is a map from value to a list of keys associated
    # with that value; 'values' is a list of values in reverse order of
    # first appearance.
    value_keys = defaultdict(list)
    values = deque()
    for k, v in elements.items():
        if v not in value_keys:
            values.appendleft(v)
        value_keys[v].append(k)

    def helper(current, r, s):
        if r == 0:
            yield current
            return
        if s == 0 or not values:
            return
        value = values.pop()
        keys = value_keys[value]
        for i in range(min(r, len(keys)), -1, -1):
            for c in combinations(keys, i):
                for result in helper(current + c, r - i, s - min(i, 1)):
                    yield result
        values.append(value)

    return helper((), r, s)

注意事项

1. 在 Python 3.3 或更高版本中，您可以使用 yield from 语句来简化递归调用:

   yield from helper(current + c, r - i, s - min(i, 1))
   

2. 如果您想知道为什么 doctests使用 collections.OrderedDict ，这样才能以可预测的顺序返回组合，这是测试工作所必需的。

3. 代码反转列表 values，并在 i 上向下迭代，这样如果调用者传入一个 OrderedDict，组合以合理的顺序返回(在输入中较早出现的值也在输出中较早出现)。

4. 考虑到从这个函数获得可预测的输出有点尴尬，我认为，值得考虑更改接口(interface)，这样就不用将键映射到值的字典，而是使用一个可迭代的 (key , 值) 对。

性能

这在速度上与 Tim Peter's combs2 大致相似:

>>> from timeit import timeit
>>> elements = dict(enumerate('abcde' * 10))
>>> test = lambda f:timeit(lambda:sum(1 for _ in f(elements, 6, 3)), number=1)
>>> test(combs2)
11.403807007009163
>>> test(constrained_combinations)
11.38378801709041