Generate all subsets of size k (containing k elements) in Python

喜夏-厌秋 提交于 2019-11-27 14:42:52

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}]

This is a subset of the power set of {1, 2, 3} (or whatever set) containing all two-element sets.

See the Python itertools documentation and search on the term "powerset" for a general answer to this problem.

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 itertools
from time import time


N = 7000
lst = [i for i in xrange(N)]

st = time()
c1 = 0
for x in itertools.combinations(lst, 2):
    c1 += 1
print "combinations: %f" % (time()-st)

st = time()
c2=0
for x in xrange(N):
    for y in xrange(x):
        c2 += 1
print "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).

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!