itertools는 iterable에서 순열/조합을 효율적으로 구현할 수 있게 도와주는 라이브러리이다.

combinations

combinations(iterable, r)

nCr : 주어진 iterable에서 순서에 상관없이 r개를 뽑아 나타낼 수 있는 조합 객체를 생성한다.

from itertools import combinations
for i in combinations([1,2,3],2):
    print(i)
    
# (1,2)
# (1,3)
# (2,3)

permutations

permutations(iterable, r)

nPr : 주어진 iterable에서 r개를 뽑아 나타낼 수 있는 순열 객체를 생성한다.

from itertools import permutations
for i in permutations([1,2,3],2):
    print(i)
    
# (1,2)
# (1,3)
# (2,1)
# (2,3)
# (3,1)
# (3,2)

product

product(*iterable, repeat=1)
n𝚷r : 복수의 iterable을 받아 중복 순열 객체를 생성한다.(repeat를 통해 같은 iterable을 여러 번 반복 가능)

from itertools import product
for i in product([1,2,3],[4,5]):
    print(i)
    
# (1,4)
# (1,5)
# (2,4)
# (2,5)
# (3,4)
# (3,5)

for j in product([1,2,3],repeat=2):
    print(j)
    
# (1,1)
# (1,2)
# (1,3)
.
.
# (2,3)
# (3,3)

combinations_with_replacement

combinations_with_replacement(iterable, r)

nHr(n+r-1Cr) : 주어진 iterable에서 중복을 허용하는 r개의 조합 객체를 생성한다.

from itertools import combinations_with_replacement as cr
for i in cr([1,2,3],2):
    print(i)
    
# (1,1)
# (1,2)
# (1,3)
# (2,2)
# (2,3)
# (3,3)

용례

프로그래머스 코딩테스트 문제 12977 : 소수 만들기에서 활용하였다.

nums 리스트를 입력으로 받아 nums에서 숫자 3개를 뽑는 경우의 수에서 소수인지를 판별

소수를 판별하는 법은 문제 제한조건에 따라 에라스토테네스의 체를 만들거나, 제곱근까지 반복문을 돌리거나 중 적절한 방법을 선택해도 무방할 것이다.

def solution(nums):
    from itertools import combinations
    
    result = 0
    
    #che = [False,False] + [True] * (3000)
    #for a in range(2,3001):
    #    for b in range(2,((3000//a)+1)):
    #        che[a*b] = False
    
    def isPrime(x):
        for i in range(2,int(x**0.5)+1):
            if x % i == 0:
                return False
        return True
    
    for i in combinations(nums,3):
        if isPrime(sum(i)):
            result += 1
    return result