최근에 학교에서 프로그래머스와 함께하는 코딩 특강을 앞두고 있습니다. 대회도 있어서 입상을 노리고 연습을 하고 있었습니다. 평소에 저는 C++로 PS를 하기 때문에 파이썬을 잘 사용하지 않아 정직하게 재귀로 백트래킹을 만들어서 정답을 받았습니다. 그런데 다른 분의 풀이를 보니 재귀 없이 itertools로 재귀를 하고 있었고 itertools에 대해 깊게 알아볼 필요성을 느껴서 자료를 정리해보게 되었습니다.
순서는 제가 중요(유용)하다고 느낀 순서대로 입니다. 활용도가 너무 낮거나 간단하게 대체할 수 있는 함수들은 임의로 제외했습니다.
itertools
python.org의 문서를 보면 itertools는 순수 파이썬에서 간결(succinct)하고 효율적인(efficient) 반복을 제공한다고 되어 있습니다. PS나 실무에서 itertools를 적절하게 활용하면 재귀를 만들거나 긴 코드 없이도 한 줄로 끝낼 수 있습니다. 매력적이지 않나요? 특히 시간이 중요한 PS에서는 많은 도움이 될 것 같습니다. 구현 실수도 없엘 수 있고요.
Python 3에서 itertools를 사용하려면 import해야 합니다. 이 글에서는 편의를 위해
from itertools import *모든 요소를 import해 두겠습니다. 이 글은 파이썬 공식 문서를 참조해 작성되었습니다.
product
주어진 iterable 들로 Cartesian Product를 수행합니다. 간단하게 생각해서 각 iterable에서 하나씩 골라서 만들 수 있는 모든 경우의 수를 만듭니다. 수식으로 표현하면 아래와 같습니다.
list(product('ABCD', 'xy'))
# [('A', 'x'),
# ('A', 'y'),
# ('B', 'x'),
# ('B', 'y'),
# ('C', 'x'),
# ('C', 'y'),
# ('D', 'x'),
# ('D', 'y')]동일한 iterable을 여러번 사용하고 싶을 때는 repeat 파라미터로 처리할 수 있습니다. 예를 들어 1, 2, 3, 4 중 하나의 값으로 7개짜리 백트레킹을 시도하고 싶으면 아래와 같이 할 수 있습니다.
len(list(product([1, 2, 3, 4], repeat=7)))
# 16384\( 4^7 = 16384 \) 경우를 모두 탐색한 것을 알 수 있습니다. product를 적절히 사용하면 n중 반복문이나 백트레킹을 편하게 처리할 수 있습니다.
permutations
입력된 iterable로 가능한 모든 순열을 만듭니다. 모든 원소를 배열해서 만들 수 있는 모든 경우를 탐색합니다. C++의 next_permutation과는 살짝 다릅니다. C++에서는 입력이 오름차순 정렬되어 있어야 모든 경우를 탐색할 수 있지만 파이썬에서는 그렇지 않습니다. 그리고 C++에서는 중복된 값이 있어서 순열 중 중복되는 것이 있으면 그건 제외됩니다. 사전순으로 가장 빠른 배열부터 늦은 배열까지 전부 찾는 느낌이라면 파이썬에서는 중복되는 값 상관 없이 배열해서 만들 수 있는 모든 경우를 찾습니다.
len(list(permutations([1, 1, 0, 0, 0])))
# 120\( 5! = 120 \)으로 중복되는 element가 있어도 모든 경우를 찾은 것을 알 수 있습니다. C++에선 미리 정렬된 배열에서 next_permutation을 통해 비슷한 작동을 얻을 수 있습니다.
#include "bits/stdc++.h"
using namespace std;
int main() {
int arr[] = { 0, 1, 1 };
do {
for (int a : arr) cout << a << ' ';
cout << '\n';
} while (next_permutation(arr, arr + 3));
}itertools의 permutation이었다면 \( 3! = 6 \)개의 경우가 나왔겠지만, 실제 출력은 011, 101, 110 단 세가지 뿐입니다. C++에선 일치하는(중복된) 경우를 제외한다는 것을 알 수 있습니다.
배열의 모든 원소를 사용하는 것이 아니라 몇개를 뽑는 경우를 찾고 싶다면 두 번째 인자로 개수를 지정해 줄 수 있습니다.
len(list(permutations([1, 1, 0, 0, 0], 2)))
# 20\( _{5}\mathrm{P}_{2} = 20 \) 개의 경우를 찾을 수 있었습니다. 마찬가지로 20개 중 중복되는 경우가 있습니다. 항상 사전순으로 빠른 경우부터 받을 수 있습니다.
combinations
\( n \)개의 요소에서 \( r \)개를 중복 없이 선택해 만들 수 있는 조합 \( _{n}\mathrm{C}_{r} \)개를 찾습니다.
list(combinations('ABCD', 2))
# [('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]마찬가지로 중복된 원소가 있으면 중복된 케이스가 있을 수 있고 항상 사전순으로 받을 수 있습니다.
combinations_with_replacement
요소에서 \( r \)개를 중복을 허용하고 선택해 만들 수 있는 조합 \( _{n}\mathrm{H}_{r} \)개를 찾습니다.
list(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')]accumulate
이름처럼 점점 숫자를 누적합니다. 입력으로 배열을 넣으면 그 배열의 누적합이 나옵니다. 누적합이 필요할 때 한 줄이면 끝낼 수 있겠네요.
a = [1, 2, 3, 4, 5]
print(list(accumulate(a)))
# [1, 3, 6, 10, 15]count
count는 특정 숫자에서 시작해 하나씩 세어나가는 함수입니다.
for i in count(10):
print(i, end=' ')
if i == 20:
break
# 10 11 12 13 14 15 16 17 18 19 20 필요에 따라 step을 2번째 인자로 제공해 증가폭을 조절할 수 있습니다.
for i in count(10, -2):
print(i, end=' ')
if i == 0:
break
# 10 8 6 4 2 0 직접 조건을 걸어 break하지 않으면 무한 루프를 돌게 됩니다. list(count(0)) 같은 코드도 위험하니 사용해선 안됩니다.
cycle
cycle은 인자로 전달된 iterable을 무한히 반복합니다.
l = 0
for c in cycle('abcd'):
print(c, end='')
l += 1
if l == 10:
break
# abcdabcdabrepeat
repeat는 첫번째 인자를 무한히 반복합니다.
l = 0
for i in repeat(10):
print(i, end=' ')
l += 1
if l == 10:
break
# 10 10 10 10 10 10 10 10 10 10두번째 인자로 반복할 횟수를 넣어 원하는 횟수만큼 반복시킬 수도 있습니다.
for i in repeat(10, 3):
print(i, end=' ')
# 10 10 10 compress
첫번째 인자로 data를 받고 두번째 인자로 selector를 받습니다. selector가 true인 data만 남겨 리턴합니다. 서로 길이가 다를 경우 둘 중 하나가 먼저 소진(exhausted)될 때까지만 반복합니다. 비트마스킹이나 필터링과 비슷한 느낌입니다.
list(compress('ABCDEF', [1, 0, 1, 0, 1, 1]))
# ['A', 'C', 'E', 'F']groupby
첫 번째 인자로 data를 받고 두 번째 인자로 함수 key를 받습니다. data의 요소를 하나씩 순회하며 key에 넣어 실행하고 리턴값이 같은 요소끼리 묶어 반환합니다. 튜플 형태로 (key, iterable) 형식으로 리턴됩니다.
for l, g in groupby(['A','B','DEF'], len):
print(l, end=' : ')
for i in g:
print(i, end=' ')
print()
# 1 : A B
# 3 : DEFbatched
첫 번째로 iterable을 주고 두 번째 인자로 정수 n을 주면 크기가 \( n \)인 배치를 만듭니다. 파이썬 3.12 버전에 추가되어서 지금 Google Colab처럼 파이썬 버전이 낮으면 사용할 수 없습니다.
data = ['roses', 'red', 'violets', 'blue', 'sugar', 'sweet']
list(batched(data, 2))
# [('roses', 'red'), ('violets', 'blue'), ('sugar', 'sweet')]\( n \)개씩 짝지어서 튜플로 반환해준다고 생각하면 됩니다.
chain
입력으로 많은 iterable을 넘겨주면 차례대로 하나씩 순회합니다.
list(chain('ABC', 'DEF', 'GHI'))
# ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I']chain.from_iterable
chain과 동일한데 입력으로 iterable을 나열하는게 아니라 배열로 전달할 수 있습니다.
list(chain.from_iterable(['ABC', 'DEF', 'GHI']))
# ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I']dropwhile
첫 번째 인자로 람다나 함수 predicate를 받고 두 번째 인자로 data를 받습니다. data의 각 요소에 대해 predicate를 수행하고 predicate가 실패(false)하면 그 원소부터 남은 원소를 전부 리턴합니다.
list(dropwhile(lambda x: x < 5, [1, 4, 6, 3, 8]))
# [6, 3, 8]6에서 람다 x < 5가 실패하기 때문에 남은 원소 전부가 출력되었습니다.
takewhile
dropwhile과 유사합니다. predicate가 계속 성공할 때까지만 리턴합니다. 한 번이라도 predicate가 실패하면 그 원소부터 뒤의 원소는 리턴되지 않습니다.
list(takewhile(lambda x: x < 5, [1, 4, 6, 3, 8]))
# [1, 4]길이를 기준으로 그룹을 만들었습니다.
zip_longest
zip은 두 iterator중 하나가 소진될 때까지 시행하는데 zip_longest는 모든 iterator가 소진될 때까지 수행합니다. 대신 한 iterator가 소진되면 None으로 채워집니다. fillvalue 파라미터를 사용해 소진되었을 때 채울 값을 지정할 수 있습니다.
list(zip('abcd', '123'))
# [('a', '1'), ('b', '2'), ('c', '3')]
list(zip_longest('abcd', '123', fillvalue='-'))
# [('a', '1'), ('b', '2'), ('c', '3'), ('d', '-')]pairwise
연속한 두 요소를 묶어 튜플로 만듭니다. 길이가 6인 배열을 파라미터로 주었으면 5개의 튜플이 나옵니다. 양 끝 원소를 제외하고 두번씩 포함됩니다.
list(pairwise('ABCDEFG'))
# [('A', 'B'), ('B', 'C'), ('C', 'D'), ('D', 'E'), ('E', 'F'), ('F', 'G')]starmap
map이랑 동일하게 동작합니다. 대신 입력이 길이가 2인 튜플인 경우에 작동합니다. 튜플로 리턴되는 함수들이랑 궁합이 좋습니다.
list(starmap(pow, [(2,5), (3,2), (10,3)]))
# [32, 9, 1000]tee
iterator를 입력으로 받아서 두 번째 인자인 n개의 iterator로 반환합니다. 각 iterator는 서로 독립(independent)입니다. 만약 서로 연결된 iterator라면 한쪽에서 순회하면 나머지 iterator가 가리키는 원소도 변하게 되는데 이걸 방지하고 싶은 경우 필요합니다.
list(tee('ABCD', 4))
# [<itertools._tee at 0x7c003f5f1540>,
# <itertools._tee at 0x7c003f5f2500>,
# <itertools._tee at 0x7c003f5f26c0>,
# <itertools._tee at 0x7c003f5f1240>]