코딩 테스트에서 도움되는 itertools 모듈

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

순열

순서가 중요할 때는 permutations를 사용한다.

서로 다른 $n$개에서 $r$개를 뽑아 순서를 고려하여 나열하는 모든 경우의 수를 구한다.
순서가 다르기 때문에 ('A', 'B')와 ('B', 'A')를 서로 다른 것으로 취급한다.

data = ["A", "B", "C"]
result = list(permutations(data, 2))  # data에서 2개를 뽑아 나열하는 모든 순열
print(result)  # [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

---

조합

combinations

순서가 상관없을 때는 combinations를 사용한다.

서로 다른 $n$개에서 $r$개를 뽑아 순서를 고려하지 않고 나열하는 모든 경우의 수를 구한다.

from itertools import combinations

result = list(combinations(data, 2))  # data에서 2개를 뽑는 모든 조합
print(result)  # [('A', 'B'), ('A', 'C'), ('B', 'C')]

combinations_with_replacement

combinations와 같지만, 중복해서 원소를 뽑는 것을 허용한다.

from itertools import combinations_with_replacement

result = list(combinations_with_replacement(data, 2))  # data에서 중복을 허용하여 2개를 뽑는 모든 조합
print(result)  # [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]

---

데카르트 곱

다중 반복문을 압축하고 싶을 때는 product를 사용한다.

여러 개의 리스트에서 각각 원소를 뽑아 조합을 만들거나, 하나의 리스트에서 중복 순열을 구할 때 사용한다.

코딩 테스트에서 복잡한 $N$중 for문을 한 줄로 깔끔하게 줄여준다.

from itertools import product

# 두 개 이상의 리스트에서 모든 조합을 구할 때
list1 = ['A', 'B']
list2 = [1, 2]
result1 = list(product(list1, list2))
print(result1)  # [('A', 1), ('A', 2), ('B', 1), ('B', 2)]

# 하나의 리스트에서 중복을 허용하여 뽑는 순열(중복 순열)을 구할 때
data = ["A", "B"]
result2 = list(product(data, repeat=2))
print(result2)  # [('A', 'A'), ('A', 'B'), ('B', 'A'), ('B', 'B')]