itertools mode 之 combinations用法

Posted 2021-02-01 phinza

itertools mode 中 combinations的功能：

from itertools import combinations
print list(combinations(list,r))
将打印 list中的所有长度为r的，子集的合集

e.g.
from itertools import combinations
list = [1,2,3]
print (list(combinations(list,2)))
>>>
[(1,2),(1,3),(2,3)]

所以当遇到排列组合的问题时，combinations会很好用

---

leetcode例题：

216. Combination Sum III

Find all possible combinations of k numbers that add up to a number n, given that only numbers from 1 to 9 can be used and each combination should be a unique set of numbers.

Note:

• All numbers will be positive integers.
• The solution set must not contain duplicate combinations.

Example 1:

Input: k = 3, n = 7
Output: [[1,2,4]]

Example 2:

Input: k = 3, n = 9
Output: [[1,2,6], [1,3,5], [2,3,4]]

---

思路：利用combinations列出符合长度k的子集，再筛选出合为n的子集
from itertools import combinations
class Solution(object):
    def combinationSum3(self, k, n):
        return [res for res in combinations(range(1,10),k) if sum(res) == n]

详细点的写法：
from itertools import combinations
class Solution(object):
    def combinationSum3(self, k, n):
        res = []
        raw = combinations(range(1,10),k)
        for i in raw:
            if sum(i) == n:
            res.append(i)
        return res