Generating All Subsets of Size k in Python

 

Suppose you are given a set \(S\) with \(n\) elements and you need to generate every subset of size \(k\) from it. For instance, if \(S=\{3,4,5\}\) and \(k=2\), then the answer would be \(\{3,4\}\), \(\{3,5\}\), \(\{4,5\}\). So, how would you do that in Python?

First of all, this is a really simple exercise. Stuff like this often comes up when someone writes a moderately large script. I wanna talk about this one in particular because it has a combinatorial solution.

Before I show it to you, I highly recommend you think about it on your own.

Solution

Notice that if you convert the set to a list, every element is assigned to a unique index between \(0\) and \(n-1\). So, it suffices to generate all the subsets of size \(k\) from \(\{0,1,\ldots,n-1\}\).

Now, every such subset falls into one of two categories: either it contains \(n-1\) or it doesn’t.

If \(n-1\) belongs to the subset, the rest of its elements form a subset of size \((k-1)\) in \(\{0,1,\ldots,n-2\}\).

If \(n-1\) doesn’t belong to the subset, then that subset is also a subset of \(\{0,1,\ldots,n-2\}\).

This is essentially a recursion. So, we are done.

Pretty neat, eh?

def gen_subset(n, k):
    if n >= 0 and n >= k >= 0:
        if n == 0 or k == 0: 
            yield set()
        elif n == k:
            # returns {0, 1, ..., n-1}
            yield set(range(n))
        else:
            # ksubsets without n-1
            yield from gen_subset(n-1, k)
            # ksubsets with n-1
            for subset in gen_subset(n-1, k-1):
                subset.add(n-1)
                yield subset
    else: 
        raise ValueError("Illegal Parameters")

 

Leave a Reply

Your email address will not be published. Required fields are marked *