22
\$\begingroup\$

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);
\$\endgroup\$
1

4 Answers 4

27
\$\begingroup\$

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.

\$\endgroup\$
6
  • 3
    \$\begingroup\$ There is a problem with that implementation / algorithm. algorithm_u(range(1,14), 3) should contain [[2, 3, 11, 12], [5, 6, 7, 8, 9], [1, 4, 10, 13]], but it doesn't. \$\endgroup\$ Jul 22, 2014 at 17:32
  • 1
    \$\begingroup\$ Confirmed with moose that something's off here. pretty_print(list(algorithm_u(["W", "X", "Y", "Z"], 2))) should contain [['W', 'X'], ['Y', 'Z']] but doesn't. \$\endgroup\$ Feb 24, 2015 at 19:04
  • \$\begingroup\$ Edited to fix the bug discovered by @moose and Brad Beattie. There were missing yield statements, as pointed out by melashot, that caused some solutions to be ignored. \$\endgroup\$ Mar 13, 2015 at 3:49
  • \$\begingroup\$ Thank you @BradBeattie for providing a smaller failing case. :) \$\endgroup\$ Mar 13, 2015 at 3:52
  • \$\begingroup\$ @melashot, I could not approve your edit in time. It helped point the way toward fixing the bug. Thanks! \$\endgroup\$ Mar 13, 2015 at 3:53
4
\$\begingroup\$

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.

\$\endgroup\$
3
\$\begingroup\$

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

\$\endgroup\$
1
\$\begingroup\$

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.

\$\endgroup\$
4
  • \$\begingroup\$ here is a timing overview between algorithm_u and the suggested subset_k method: i.stack.imgur.com/nQt5B.png \$\endgroup\$
    – eltings
    Apr 13, 2022 at 13:57
  • \$\begingroup\$ Nice to see the results. I would imagine algorithm_u might perform better for very large n and k, obviously the simpler algorithms always perform better up to some constant usually, depending on implementation details and such. Though it is unclear in this case. \$\endgroup\$ Apr 14, 2022 at 14:19
  • \$\begingroup\$ Is that the fastest version of all those answers? for N around 15, it's quite slow. If yes, I'm wondering how much performance would be gained for this to be written in C or C++ \$\endgroup\$
    – jokoon
    Aug 26, 2022 at 8:57
  • \$\begingroup\$ Also curious if pypy would help. \$\endgroup\$
    – jokoon
    Aug 26, 2022 at 12:20

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

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