生成所有固定长度组合的组

5 投票
1 回答
1382 浏览
提问于 2025-04-17 13:55

我在寻找一个算法或者Python代码,用来生成将一组有n个元素分成零个或多个包含r个元素的组以及剩余元素的所有可能方式。比如,给定一个集合:

[1,2,3,4,5]

n = 5r = 2时,我想得到:

((1,2,3,4,5),)
((1,2),(3,4,5))
((1,3),(2,4,5))
...
((1,2),(3,4),(5,))
((1,2),(3,5),(4,))
...

换句话说,就是从这个集合中提取0组两个元素的结果,加上提取1组两个元素的结果,再加上提取2组两个元素的结果……如果n更大,这个过程会继续下去。

这些结果生成的顺序并不重要,每个组内元素的顺序也无所谓,组与组之间的顺序同样不重要。(例如,((1,3),(2,4,5))((3,1),(4,5,2))以及((2,5,4),(1,3))都是等价的。)我希望每个不同的结果至少能生成一次,最好是刚好生成一次,并且尽可能高效。


最简单的方法是生成从n个元素中选择r个元素的所有可能组合,然后创建这些组合的所有可能组(也就是幂集),遍历这些组,只处理那些组合之间没有共同元素的组。可是,这种方法对于即使是少量元素来说也太慢了(它需要遍历2^(n!/r!(n-r)!)个组,所以复杂度是双指数的)。

基于这个问题中给出的代码,实际上是r = 2n为偶数的特例,我想出了以下代码:

def distinct_combination_groups(iterable, r):
    tpl = tuple(iterable)
    yield (tpl,)
    if len(tpl) > r:
        for c in combinations(tpl, r):
            for g in distinct_combination_groups(set(tpl) - set(c), r):
                yield ((c,) + g)

这段代码似乎能生成所有可能的结果,但当n比较大时,它会包含一些重复的结果,数量还不少。所以我希望能找到一个避免重复的算法。

1 个回答

9

这个怎么样呢?

from itertools import combinations

def partitions(s, r):
    """
    Generate partitions of the iterable `s` into subsets of size `r`.

    >>> list(partitions(set(range(4)), 2))
    [((0, 1), (2, 3)), ((0, 2), (1, 3)), ((0, 3), (1, 2))]
    """
    s = set(s)
    assert(len(s) % r == 0)
    if len(s) == 0:
        yield ()
        return
    first = next(iter(s))
    rest = s.difference((first,))
    for c in combinations(rest, r - 1):
        first_subset = (first,) + c
        for p in partitions(rest.difference(c), r):
            yield (first_subset,) + p

def partitions_with_remainder(s, r):
    """
    Generate partitions of the iterable `s` into subsets of size
    `r` plus a remainder.

    >>> list(partitions_with_remainder(range(4), 2))
    [((0, 1, 2, 3),), ((0, 1), (2, 3)), ((0, 2), (1, 3)), ((0, 3), (1, 2))]
    >>> list(partitions_with_remainder(range(3), 2))
    [((0, 1, 2),), ((1, 2), (0,)), ((0, 2), (1,)), ((0, 1), (2,))]
    """
    s = set(s)
    for n in xrange(len(s), -1, -r): # n is size of remainder.
        if n == 0:
            for p in partitions(s, r):
                yield p
        elif n != r:
            for remainder in combinations(s, n):
                for p in partitions(s.difference(remainder), r):
                    yield p + (remainder,)

这是提问者给出的例子:

>>> pprint(list(partitions_with_remainder(range(1, 6), 2)))
[((1, 2, 3, 4, 5),),
 ((4, 5), (1, 2, 3)),
 ((3, 5), (1, 2, 4)),
 ((3, 4), (1, 2, 5)),
 ((2, 5), (1, 3, 4)),
 ((2, 4), (1, 3, 5)),
 ((2, 3), (1, 4, 5)),
 ((1, 5), (2, 3, 4)),
 ((1, 4), (2, 3, 5)),
 ((1, 3), (2, 4, 5)),
 ((1, 2), (3, 4, 5)),
 ((2, 3), (4, 5), (1,)),
 ((2, 4), (3, 5), (1,)),
 ((2, 5), (3, 4), (1,)),
 ((1, 3), (4, 5), (2,)),
 ((1, 4), (3, 5), (2,)),
 ((1, 5), (3, 4), (2,)),
 ((1, 2), (4, 5), (3,)),
 ((1, 4), (2, 5), (3,)),
 ((1, 5), (2, 4), (3,)),
 ((1, 2), (3, 5), (4,)),
 ((1, 3), (2, 5), (4,)),
 ((1, 5), (2, 3), (4,)),
 ((1, 2), (3, 4), (5,)),
 ((1, 3), (2, 4), (5,)),
 ((1, 4), (2, 3), (5,))]

撰写回答