itertools

이 포스트는 2021.12~2022.09 기간동안 벨로그에 작성한 글을 티스토리에 옮겨 적은 것입니다.

코테 연습 중 리스트에서 3개의 원소를 조합 중복 없이 뽑아야 하는 문제가 있어서 3중루프를 돌며 뻘짓을 하고 있었다.🤯 그런데 알고 보니 itertools 라이브러리를 사용하면 한 줄로 해결할 수 있었다....🤦‍♀️

itertools

파이썬에서 반복되는 데이터를 처리하는 기능을 포함하고 있는 라이브러리이다.

permutations
리스트와 같은 iterable 객체에서 r개의 데이터를 뽑아 일렬로 나열하는 모든 경우(순열=permutations)을 계산한다. 시간복잡도는 O(n!)이다.

from itertools import permutations

data=['A','B','C']
#data 리스트에서 3개 데이터를 뽑아 모든 경우를 나열한 후 리스트로 변환.
result=list(permutations(data,3))
print(result)
'''
실행결과
[('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'), 
('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]
'''

combinations
리스트와 같은 iterable 객체에서 r개의 데이터를 뽑아 순서를 고려하지 않고 나열하는 모든 경우(조합=combinations)을 계산한다.

from itertools import combinations

data=['A','B','C']
#data 리스트에서 2개 데이터를 뽑아 모든 조합을 구하고 리스트로 변환
result=list(combinations(data,2))
print(result)
'''
실행결과
[('A', 'B'), ('A', 'C'), ('B', 'C')]
'''

product
product는 permutations과 같지만 원소 중복을 허용한다. 뽑고자 하는 데이터의 수를 repeat 속성으로 넣는다.

from itertools import product

data=['A','B','C']

result=list(product(data,repeat=2))
print(result)
'''
실행결과
[('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'A'), 
('B', 'B'), ('B', 'C'), ('C', 'A'), ('C', 'B'), ('C', 'C')]
'''

combinations_with_replacement
combinations_with_replacement는 combinations와 같지만 원소 중복을 허용한다.

from itertools import combinations_with_replacement

data=['A','B','C']
 
result=list(combinations_with_replacement(data,2))
print(result)
'''
실행결과
[('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]
'''

※map을 활용할 수도 있다.

res=list(map(list,permutations(s,2)))
res=list(map(커스텀함수,permutations(s,2)))

프로그래머스 고득점kit 완전탐색 부분 소수찾기 문제 참고

참고자료: 이코테 452p~454p