Take the 2-minute tour ×
Code Review Stack Exchange is a question and answer site for peer programmer code reviews. It's 100% free, no registration required.

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);
share|improve this question
add comment

4 Answers 4

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:
                visit(n, a)
                a[nu] = a[nu] + 1
            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:
            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; 12|3|45; 1|23|45; 13|2|45; 1|2|345; 1|24|35; 14|2|35; 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.

share|improve this answer
add comment

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.

share|improve this answer
add comment

Try this one:

def k_subset(k, list_):
    min_per_set = len(list_) / k
    remaining = len(list_) % k
    final_sets = []
    buffer_ = list_[:]

    for i in range(0, k):
        if remaining > 0:
            in_this_set = min_per_set + 1
            remaining -= 1
        else:
            in_this_set = min_per_set
        this_set = buffer_[:in_this_set]
        buffer_ = buffer_[in_this_set:len(buffer_)]
        final_sets.append(this_set)

    return final_sets


print k_subset(3, [1, 2, 3, 4])
share|improve this answer
1  
This generates one k_subset. I'd like to generate all of them –  spinning_plate Mar 29 '11 at 22:54
add comment

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

share|improve this answer
add comment

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.