如何使用itertools在Python中找到最小长度为N的元素组合

Question

from itertools import combinations
a = [1,2,3]
combinations(a,2) #will give me ((1,2),(1,3),(2,3))
combinations(a,3) #will give me ((1,2,3),)

但是如果我想要一个数组中不同长度的结果呢？

比如说，我想找到数组[1,2,3]中长度大于或等于2的所有组合，结果应该是 ((1,2),(1,3),(2,3),(1,2,3))

类似于 c = combinations(a,>=2) 这样的用法。

我试着用lambda函数，但没有成功。

c = combinations(a,lambda x: x for x in [2,3])

还有用列表推导式 c = combinations(a,[x for x in [2,3]])

我知道我可以用一个简单的循环来找出不同长度的组合。

for l in [2,3]:
    combinations(a,l)

但是有没有什么更“pythonic”的方法来做到这一点呢？

Answer 1

你可以把 combinations 和 chain.from_iterable 结合起来使用：

>>> from itertools import chain, combinations
>>> a = [1,2,3]
>>> n = 2
>>> cc = chain.from_iterable(combinations(a, i) for i in range(n, len(a)+1))
>>> list(cc)
[(1, 2), (1, 3), (2, 3), (1, 2, 3)]

这里的 chain.from_iterable 是在把生成器表达式 (combinations(a, i) for i in range(n, len(a)+1)) 产生的结果“压扁”，也就是把里面的内容整理成一行。否则你得到的结果可能会像这样：

>>> [list(combinations(a,i)) for i in range(n, len(a)+1)]
[[(1, 2), (1, 3), (2, 3)], [(1, 2, 3)]]

这样也可以，但格式上可能不是你想要的那样。