코딩 테스트를 위한 파이썬 문법 - itertools (책: 이것이 코딩 테스트다)

itertools

permutations

리스트와 같은 iterable 객체에서 r개의 데이터를 뽑아 일렬로 나열하는 모든 경우(순열)를 계산해준다.

클래스이므로 객체 초기화 후 리스트 자료형으로 변환하여 사용한다.

예시: 대회 순위

from itertools import permutations

data = ['A', 'B', 'C']
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개의 데이터를 뽑아 순서를 고려하지 않고 나열하는 모든 경우(조합)를 계산한다.

클래스이므로 객체 초기화 후 리스트 자료형으로 변환하여 사용한다.

예시: 로또

from itertools import combinations

data = ['A', 'B', 'C']
result = list(combinations(data, 3))

print(result)

[('A', 'B'), ('A', 'C'), ('B', 'C')]

product

permutations와 같이 iterable 객체에서 r개의 데이터를 뽑아 일렬로 나열하는 모든 경우(순열)를 계산한다. 다만 원소를 중복하여 뽑는다.

뽑고자 하는 데이터의 수를 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와 같이 리스트와 같은 iterable 객체에서 r개의 데이터를 뽑아 순서를 고려하지 않고 나열하는 모든 경우(조합)를 계산한다. 다만 원소를 중복해서 뽑는다.

클래스이므로 객체 초기화 후 리스트 자료형으로 변환하여 사용한다.

예시: 아이스크림 맛 선택

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')]

요약

함수 정의 현실 예시

combinations	순서 상관없이 n개 선택	로또 번호 선택
permutations	순서를 고려하여 n개 선택	달리기 대회 순위
product	중복 허용하여 모든 조합	비밀번호 brute-force 시도
combinations_with_replacement	중복 허용하며 순서 상관없이 n개 선택	아이스크림 맛 선택