Question

I’m looking for an algorithm and/or Python code to generate all possible ways of partitioning a set of n elements into zero or more groups of r elements and a remainder. For example, given a set:

[1,2,3,4,5]

with n = 5 and r = 2, I would like to get

((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,))
...

In other words, the result of extracting 0 groups of two items from the set, plus the results of extracting 1 group of two items from the set, plus the results of extracting 2 groups of two from the set,… if n were larger, this would continue.

The order in which these results are generated is not important, and neither is the order of elements within each individual group, nor the order of the groups within a result. (e.g. ((1,3),(2,4,5)) is equivalent to ((3,1),(4,5,2)) and to ((2,5,4),(1,3)) and so on.) What I’m looking for is that every distinct result is produced at least once, and preferably exactly once, in as efficient a manner as possible.

---

The brute force method is to generate all possible combinations of r out of the n elements, then create all possible groups of any number of those combinations (the powerset), iterate over them and only process the ones where the combinations in the group have no elements in common. That takes far too long for even a small number of elements (it requires iterating over 2^(n!/r!(n-r)!) groups, so the complexity is double-exponential).

Based on the code given in this question, which is essentially the special case for r = 2 and n even, I’ve come up with the following:

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)

which does seem to generate all possible results, but it includes some duplicates, a nontrivial number of them when n is fairly large. So I’m hoping for an algorithm that will avoid the duplicates.

Answer 1

How about this?

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,)

The example from the OP:

>>> 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,))]