Generate all subsets of size k (containing k elements) in Python
Seems like you want itertools.combinations
:
>>> list(itertools.combinations((1, 2, 3), 2))[(1, 2), (1, 3), (2, 3)]
If you want sets you'll have to convert them explicitly. If you don't mind an iterable instead of a list, and you're using Python 3, you can use map
:
>>> s = set((1, 2, 3))>>> map(set, itertools.combinations(s, 2))<map object at 0x10cdc26d8>
To view all the results at once, you can pass the output of map
to list
. (In Python 2, the output of map
is automatically a list.)
>>> list(map(set, itertools.combinations(s, 2)))[{1, 2}, {1, 3}, {2, 3}]
However, if you know you'll need a list, a list comprehension is marginally better (h/t Jacob Bowyer):
>>> [set(i) for i in itertools.combinations(s, 2)][{1, 2}, {1, 3}, {2, 3}]
Just to give another perspective, I looked for a way to iterate all subset of size 2 of {1.....N}
, so I put itertools.combinations
into test:
import itertoolsfrom time import timeN = 7000lst = [i for i in xrange(N)]st = time()c1 = 0for x in itertools.combinations(lst, 2): c1 += 1print "combinations: %f" % (time()-st)st = time()c2=0for x in xrange(N): for y in xrange(x): c2 += 1print "double loop: %f" % (time()-st)print "c1=%d,c2=%d" % (c1,c2)# prints:#combinations: 4.247000#double loop: 3.479000# c1=24496500,c2=24496500
So I guess you should not always turn into the general solution.... If you know in advance the size of the subset you want, it should be more efficient to iterate using for loops.
Also note that you should not iterate over list(itertools.combinations(lst, 2))
since this move creates the list (and much slower than using the generator itself).