[파이썬] itertools모듈

리스트 조합 관련된 알고리즘 문제를 풀다 보면

itertools 라이브러리가 풀이에 종종 등장한다.

 

사실 처음 봤을 땐

라이브러리 사용법을 암기해서 쓰는 것보다

직접 구현하는 게 차라리 빠르겠다 싶었는데,

생각보다 사용빈도나 범위가 큰 것 같아서 정리해두려고 한다.

 

itertools는 Python 내장 라이브러리로, 

자신만의 반복자를 만들 수 있는 모듈이다.

 

주로 사용되는 함수들은 다음과 같다.

 

● product()

함수 이름 그대로 곱집합(Cartesian product)을 만드는 함수이다.

iterable객체를 반환한다.

(편의상 출력하는 과정은 생략하였다. 동일하게 결과를 얻으려면 list로 변환한 뒤 print함수로 출력하면 된다.)

import itertools
itertools.product('ABCD', repeat=2)
# 결과 : [('A', 'A'), ('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'A'), ('B', 'B'), ('B', 'C'), ('B', 'D'), ('C', 'A'), ('C', 'B'), ('C', 'C'), ('C', 'D'), ('D', 'A'), ('D', 'B'), ('D', 'C'), ('D', 'D')]

import itertools
list(itertools.product('ABCD', 'XY', '1234'))
# 결과 : [('A', 'X', '1'), ('A', 'X', '2'), ('A', 'X', '3'), ('A', 'X', '4'), ('A', 'Y', '1'), ('A', 'Y', '2'), ('A', 'Y', '3'), ('A', 'Y', '4'), ('B', 'X', '1'), ('B', 'X', '2'), ('B', 'X', '3'), ('B', 'X', '4'), ('B', 'Y', '1'), ('B', 'Y', '2'), ('B', 'Y', '3'), ('B', 'Y', '4'), ('C', 'X', '1'), ('C', 'X', '2'), ('C', 'X', '3'), ('C', 'X', '4'), ('C', 'Y', '1'), ('C', 'Y', '2'), ('C', 'Y', '3'), ('C', 'Y', '4'), ('D', 'X', '1'), ('D', 'X', '2'), ('D', 'X', '3'), ('D', 'X', '4'), ('D', 'Y', '1'), ('D', 'Y', '2'), ('D', 'Y', '3'), ('D', 'Y', '4')]

 

● permutation()

길이가 r인 순열을 만드는 함수이다.

import itertools
itertools.permutations('ABCD', 2)
# 결과 : [('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'A'), ('B', 'C'), ('B', 'D'), ('C', 'A'), ('C', 'B'), ('C', 'D'), ('D', 'A'), ('D', 'B'), ('D', 'C')]

수학에서의 순열과 동일하게, 4P2의 결과라고 보면 된다.

 

● combinations()

길이가 r인 조합을 만드는 함수이다.

import itertools
itertools.combinations('ABCD', 2)
# 결과 : [('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]

4C2의 결과라고 보면 된다.

 

● combinations_with_replacement()

위의 combinations()와 동일하나 반복 사용을 허용한다.

그렇기 때문에 인자로 받은 iterable의 원소의 개수보다 r이 커도 된다.

import itertools
itertools.combinations_with_replacement('ABCD', 2)
# 결과 : [('A', 'A'), ('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'B'), ('B', 'C'), ('B', 'D'), ('C', 'C'), ('C', 'D'), ('D', 'D')]

4H2의 결과라고 보면 된다.

---

공식문서:

itertools — 효율적인 루핑을 위한 이터레이터를 만드는 함수 — Python 3.9.1 문서