每个元素的最大重复次数的组合

Question

我想获得一个 k 大小的元组列表，其中包含一个元素列表的组合（我们称之为elements），类似于itertools.combinations_with_replacement(elements, k) 所做的事情。不同之处在于我想为每个元素的替换数量添加一个最大值。

例如，如果我运行以下命令：

elements = ['a', 'b']
print(list(itertools.combinations_with_replacement(elements, 3)))

我明白了：

[('a', 'a', 'a'), ('a', 'a', 'b'), ('a', 'b', 'b'), ('b', 'b', 'b')]

我想要以下内容：

elements = {'a': 2, 'b': 3}
print(list(combinations_with_max_replacement(elements, 3)))

哪个会打印

[('a', 'a', 'b'), ('a', 'b', 'b'), ('b', 'b', 'b')]

请注意，每个元组中'a' 的最大数量是2，因此('a', 'a', 'a') 不是结果的一部分。

我宁愿避免循环遍历 itertools.combinations_with_replacement(elements, k) 计算每个元组中的元素并将它们过滤掉的结果。

如果我可以提供更多信息，请告诉我。

感谢您的帮助！

更新

我试过了：

elements = ['a'] * 2 + ['b'] * 3
print(set(itertools.combinations(elements, 3)))

然后得到：

{('a', 'b', 'b'), ('b', 'b', 'b'), ('a', 'a', 'b')}

我得到了我需要的元素，但我失去了顺序，看起来有点老套

Answer 1

我知道您不想循环遍历结果，但以这种方式过滤输出可能更容易。

def custom_combinations(elements, max_count):
    L = list(itertools.combinations_with_replacement(elements, max_count))
    for element in elements.keys():
        L = list(filter(lambda x: x.count(element) <= elements[element], L))
    return L

Answer 2

纯 Python 解决方案（即没有itertools）

你可以使用递归：

def combos(els, l):
    if l == 1:
        return [(k,) for k, v in els.items() if v]
    cs = []
    for e in els:
        nd = {k: v if k != e else v - 1 for k, v in els.items() if v}
        cs += [(e,)+c for c in combos(nd, l-1)]
    return cs

测试表明它有效：

>>> combos({'a': 2, 'b': 3}, 3)
[('b', 'b', 'b'), ('b', 'b', 'a'), ('b', 'a', 'b'), ('b', 'a', 'a'), ('a', 'b', 'b'), ('a', 'b', 'a'), ('a', 'a', 'b')]

请注意，我们确实会取消订单，但如果我们按照您的要求将 els 作为 dictionary 传递，这是不可避免的。

Answer 3

我相信这种递归解决方案具有您想要的时间复杂度。

我们不是传递一个字典，而是传递一个项目对的列表。我们还传递了start_idx，它告诉“较低”递归函数调用忽略较早的元素。这解决了另一个递归答案的乱序问题。

def _combos(elements, start_idx, length):
    # ignore elements before start_idx
    for i in range(start_idx, len(elements)):
        elem, count = elements[i]
        if count == 0:
            continue
        # base case: only one element needed
        if length == 1:
            yield (elem,)
        else:
            # need more than one elem: mutate the list and recurse
            elements[i] = (elem, count - 1)
            # when we recurse, we ignore elements before this one
            # this ensures we find combinations, not permutations
            for combo in _combos(elements, i, length - 1):
                yield (elem,) + combo
            # fix the list
            elements[i] = (elem, count)


def combos(elements, length):
    elements = list(elements.items())
    return _combos(elements, 0, length)

print(list(combos({'a': 2, 'b': 3}, 3)))
# [('a', 'a', 'b'), ('a', 'b', 'b'), ('b', 'b', 'b')]

作为奖励，随着输入大小的增长，分析表明它比 set(itertools.combinations(_)) 解决方案的性能更高。

print(timeit.Timer("list(combos({'a': 2, 'b': 2, 'c': 2}, 3))",
             setup="from __main__ import combos").timeit())
# 9.647649317979813
print(timeit.Timer("set(itertools.combinations(['a'] * 2 + ['b'] * 2 + ['c'] * 2, 3))").timeit())
# 1.7750148189952597

print(timeit.Timer("list(combos({'a': 4, 'b': 4, 'c': 4}, 4))",
             setup="from __main__ import combos").timeit())
# 20.669851204031147
print(timeit.Timer("set(itertools.combinations(['a'] * 4 + ['b'] * 4 + ['c'] * 4, 4))").timeit())
# 28.194088937016204

print(timeit.Timer("list(combos({'a': 5, 'b': 5, 'c': 5}, 5))",
             setup="from __main__ import combos").timeit())
# 36.4631432640017
print(timeit.Timer("set(itertools.combinations(['a'] * 5 + ['b'] * 5 + ['c'] * 5, 5))").timeit())
# 177.29063899395987