Finding all k-subset partitions

Question

The following code generates all \$k\$-subsets of a given array. A \$k\$-subset of set \$X\$ is a partition of all the elements in \$X\$ into \$k\$ non-empty subsets.

Thus, for {1,2,3,4} a 3-subset is {{1,2},{3},{4}}.

I'm looking for improvements to the algorithm or code. Specifically, is there a better way than using copy.deepcopy? Is there some itertools magic that does this already?

import copy
arr = [1,2,3,4]

def t(k,accum,index):
    print accum,k
    if index == len(arr):
        if(k==0):
            return accum;
        else:
            return [];

    element = arr[index];
    result = []

    for set_i in range(len(accum)):
        if k>0:
            clone_new = copy.deepcopy(accum);
            clone_new[set_i].append([element]);
            result.extend( t(k-1,clone_new,index+1) );

        for elem_i in range(len(accum[set_i])):
            clone_new = copy.deepcopy(accum);
            clone_new[set_i][elem_i].append(element)
            result.extend( t(k,clone_new,index+1) );

    return result

print t(3,[[]],0);

Answer 1

A very efficient algorithm (Algorithm U) is described by Knuth in the Art of Computer Programming, Volume 4, Fascicle 3B to find all set partitions with a given number of blocks. Your algorithm, although simple to express, is essentially a brute-force tree search, which is not efficient.

Since Knuth's algorithm is not very concise, its implementation is lengthy as well. Note that the implementation below moves an item among the blocks one at a time and need not maintain an accumulator containing all partial results. For this reason, no copying is required.

def algorithm_u(ns, m):
    def visit(n, a):
        ps = [[] for i in xrange(m)]
        for j in xrange(n):
            ps[a[j + 1]].append(ns[j])
        return ps

    def f(mu, nu, sigma, n, a):
        if mu == 2:
            yield visit(n, a)
        else:
            for v in f(mu - 1, nu - 1, (mu + sigma) % 2, n, a):
                yield v
        if nu == mu + 1:
            a[mu] = mu - 1
            yield visit(n, a)
            while a[nu] > 0:
                a[nu] = a[nu] - 1
                yield visit(n, a)
        elif nu > mu + 1:
            if (mu + sigma) % 2 == 1:
                a[nu - 1] = mu - 1
            else:
                a[mu] = mu - 1
            if (a[nu] + sigma) % 2 == 1:
                for v in b(mu, nu - 1, 0, n, a):
                    yield v
            else:
                for v in f(mu, nu - 1, 0, n, a):
                    yield v
            while a[nu] > 0:
                a[nu] = a[nu] - 1
                if (a[nu] + sigma) % 2 == 1:
                    for v in b(mu, nu - 1, 0, n, a):
                        yield v
                else:
                    for v in f(mu, nu - 1, 0, n, a):
                        yield v

    def b(mu, nu, sigma, n, a):
        if nu == mu + 1:
            while a[nu] < mu - 1:
                yield visit(n, a)
                a[nu] = a[nu] + 1
            yield visit(n, a)
            a[mu] = 0
        elif nu > mu + 1:
            if (a[nu] + sigma) % 2 == 1:
                for v in f(mu, nu - 1, 0, n, a):
                    yield v
            else:
                for v in b(mu, nu - 1, 0, n, a):
                    yield v
            while a[nu] < mu - 1:
                a[nu] = a[nu] + 1
                if (a[nu] + sigma) % 2 == 1:
                    for v in f(mu, nu - 1, 0, n, a):
                        yield v
                else:
                    for v in b(mu, nu - 1, 0, n, a):
                        yield v
            if (mu + sigma) % 2 == 1:
                a[nu - 1] = 0
            else:
                a[mu] = 0
        if mu == 2:
            yield visit(n, a)
        else:
            for v in b(mu - 1, nu - 1, (mu + sigma) % 2, n, a):
                yield v

    n = len(ns)
    a = [0] * (n + 1)
    for j in xrange(1, m + 1):
        a[n - m + j] = j - 1
    return f(m, n, 0, n, a)

Examples:

def pretty_print(parts):
    print '; '.join('|'.join(''.join(str(e) for e in loe) for loe in part) for part in parts)

>>> pretty_print(algorithm_u([1, 2, 3, 4], 3))
12|3|4; 1|23|4; 13|2|4; 1|2|34; 1|24|3; 14|2|3

>>> pretty_print(algorithm_u([1, 2, 3, 4, 5], 3))
123|4|5; 12|34|5; 1|234|5; 13|24|5; 134|2|5; 14|23|5; 124|3|5; 12|3|45; 1|23|45; 13|2|45; 1|2|345; 1|24|35; 14|2|35; 14|25|3; 1|245|3; 1|25|34; 13|25|4; 1|235|4; 12|35|4; 125|3|4; 15|23|4; 135|2|4; 15|2|34; 15|24|3; 145|2|3

Timing results:

$ python -m timeit "import test" "test.t(3, [[]], 0, [1, 2, 3, 4])"
100 loops, best of 3: 2.09 msec per loop

$ python -m timeit "import test" "test.t(3, [[]], 0, [1, 2, 3, 4, 5])"
100 loops, best of 3: 7.88 msec per loop

$ python -m timeit "import test" "test.t(3, [[]], 0, [1, 2, 3, 4, 5, 6])"
10 loops, best of 3: 23.6 msec per loop

$ python -m timeit "import test" "test.algorithm_u([1, 2, 3, 4], 3)"
10000 loops, best of 3: 26.1 usec per loop

$ python -m timeit "import test" "test.algorithm_u([1, 2, 3, 4, 5, 6, 7, 8], 3)"
10000 loops, best of 3: 28.1 usec per loop

$ python -m timeit "import test" "test.algorithm_u([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], 3)"
10000 loops, best of 3: 29.4 usec per loop

Notice that t runs much slower than algorithm_u for the same input. Furthermore, t runs exponentially slower with each extra input, whereas algorithm_u runs almost as fast for double and quadruple the input size.

Answer 2

Here's the obligatory recursive version:

def k_subset(s, k):
    if k == len(s):
        return (tuple([(x,) for x in s]),)
    k_subs = []
    for i in range(len(s)):
        partials = k_subset(s[:i] + s[i + 1:], k)
        for partial in partials:
            for p in range(len(partial)):
                k_subs.append(partial[:p] + (partial[p] + (s[i],),) + partial[p + 1:])
    return k_subs

This returns a bunch of duplicates which can be removed using

def uniq_subsets(s):
    u = set()
    for x in s:
        t = []
        for y in x:
            y = list(y)
            y.sort()
            t.append(tuple(y))
        t.sort()
        u.add(tuple(t))
    return u

So the final product can be had with

print uniq_subsets(k_subset([1, 2, 3, 4], 3))

set([
    ((1,), (2,), (3, 4)), 
    ((1,), (2, 4), (3,)), 
    ((1, 3), (2,), (4,)), 
    ((1, 4), (2,), (3,)), 
    ((1,), (2, 3), (4,)), 
    ((1, 2), (3,), (4,))
])

Wow, that's pretty bad and quite unpythonic. :(

Edit: Yes, I realize that reimplementing the problem doesn't help with reviewing the original solution. I was hoping to gain some insight on your solution by doing so. If it's utterly unhelpful, downvote and I'll remove the answer.

Edit 2: I removed the unnecessary second recursive call. It's shorter but still not very elegant.

Answer 3

Couldn't find any super quick wins in itertools.

(Maybe I didn't look hard enough.) I did however come up with this, it runs pretty slow, but is fairly concise:

from itertools import chain, combinations

def subsets(arr):
    """ Note this only returns non empty subsets of arr"""
    return chain(*[combinations(arr,i + 1) for i,a in enumerate(arr)])

def k_subset(arr, k):
    s_arr = sorted(arr)
    return set([frozenset(i) for i in combinations(subsets(arr),k) 
               if sorted(chain(*i)) == s_arr])


print k_subset([1,2,3,4],3)

Some minor wins for speed but less concise, would be to only do the set//frozenset buisness at the end if there are non unique elements in the array, or use a custom flatten function or sum(a_list, []) rather than chain(*a_list).

If you are desperate for speed you might want have a think about another language or maybe: www.cython.org is pretty neat. Obviously the above algorithms are much better speed-wise for starters.

Also what might be worht a look into is www.sagemath.org... It's an mathematical environment built atop of python, and for example functions like, subsets() , flatten etc and lots of combinatorial things live there.

Cheers,

Matt

Answer 4

Based on an answer in an all partitions question (https://stackoverflow.com/questions/19368375/set-partitions-in-python/61141601): This can be done with simple recursion, no need for itertools, no need for a complicated algorithm. I am surprised this was not suggested. It should be just as efficient as Knuth's algorithm as well as it goes through every combination only once.

def subsets_k(collection, k): yield from partition_k(collection, k, k)
def partition_k(collection, min, k):
  if len(collection) == 1:
    yield [ collection ]
    return

  first = collection[0]
  for smaller in partition_k(collection[1:], min - 1, k):
    if len(smaller) > k: continue
    # insert `first` in each of the subpartition's subsets
    if len(smaller) >= min:
      for n, subset in enumerate(smaller):
        yield smaller[:n] + [[ first ] + subset]  + smaller[n+1:]
    # put `first` in its own subset 
    if len(smaller) < k: yield [ [ first ] ] + smaller

Test:

something = list(range(1,5))
for n, p in enumerate(subsets_k(something, 1), 1):
  print(n, sorted(p))
for n, p in enumerate(subsets_k(something, 2), 1):
  print(n, sorted(p))
for n, p in enumerate(subsets_k(something, 3), 1):
  print(n, sorted(p))  
for n, p in enumerate(subsets_k(something, 4), 1):
  print(n, sorted(p))

Yields correctly:

1 [[1, 2, 3, 4]]

1 [[1], [2, 3, 4]]
2 [[1, 2], [3, 4]]
3 [[1, 3, 4], [2]]
4 [[1, 2, 3], [4]]
5 [[1, 4], [2, 3]]
6 [[1, 3], [2, 4]]
7 [[1, 2, 4], [3]]

1 [[1], [2], [3, 4]]
2 [[1], [2, 3], [4]]
3 [[1], [2, 4], [3]]
4 [[1, 2], [3], [4]]
5 [[1, 3], [2], [4]]
6 [[1, 4], [2], [3]]

1 [[1], [2], [3], [4]]

Compared to the Knuth implementation above (modified for Python3 - xrange changed to range):

if __name__ == '__main__':
  import timeit
  print(timeit.timeit("for _ in subsets_k([1, 2, 3, 4, 5], 3): pass", globals=globals()))
  print(timeit.timeit("for _ in algorithm_u([1, 2, 3, 4, 5], 3): pass", globals=globals()))

Results in more than twice as fast code:

20.724652599994442
41.03094519999286

Sometimes a simple approach is the best one. It would be interesting to know if optimizations can be applied to Knuth variation to fix this, or if this simple algorithm is the best one.

Update: the Knuth timing information above is both wrong and misleading!!!

The t implementation compiles a whole list and does not return a generator. Whereas the Knuth version has a generator. To make a fair test comparison, one must enumerate all the elements otherwise the Knuth implementation is just running to the first yield returning the generator and timing this... for _ in generator: pass would have been sufficient to have a real test comparison.