[Python] itertools: 효율적인 iterator 빌딩 블록
permutation, python itertools, chain, cycle, islice, groupby, combinations, permutations, product
정의
itertools는 메모리 효율적인 iterator 조합기를 제공한다. lazy 평가로 무한 시퀀스도 다루며, C로 구현되어 빠르다. Haskell/Lisp 영향을 받은 함수형 도구함.
무한 iterator
count
import itertools
for i in itertools.count(start=10, step=2):
if i > 20: break
print(i)
# 10 12 14 16 18 20
cycle
for c in itertools.cycle("ABC"):
print(c) # A B C A B C A B C ... (무한)
repeat
list(itertools.repeat("x", 3)) # ['x', 'x', 'x']
# zip과 결합해 인덱스 생성
list(zip(itertools.count(), "abc")) # [(0, 'a'), (1, 'b'), (2, 'c')]
# 보통 enumerate 사용이 더 Pythonic
iterator 단축/슬라이싱
islice
iterator의 슬라이스. generator/무한 시퀀스에도 적용 가능 (일반 슬라이스는 시퀀스만).
list(itertools.islice(range(100), 5)) # [0, 1, 2, 3, 4]
list(itertools.islice(range(100), 5, 10)) # [5, 6, 7, 8, 9]
list(itertools.islice(range(100), 0, 10, 2)) # [0, 2, 4, 6, 8]
takewhile / dropwhile
조건 만족하는 동안 / 안 만족할 때까지.
python
import itertools
xs = [1, 3, 5, 7, 4, 8, 2]
print(list(itertools.takewhile(lambda x: x < 6, xs)))
print(list(itertools.dropwhile(lambda x: x < 6, xs)))
# filter와 다름: filter는 모든 원소 확인, takewhile은 첫 falsy에서 중단
print(list(filter(lambda x: x < 6, xs))) 결과
[1, 3, 5]
[7, 4, 8, 2]
[1, 3, 5, 4, 2]결합
chain
여러 iterable을 하나로 평탄화.
list(itertools.chain([1, 2], (3, 4), {5, 6}))
# [1, 2, 3, 4, 5, 6]
# chain.from_iterable: 이터러블의 이터러블
list(itertools.chain.from_iterable([[1, 2], [3, 4], [5]]))
# [1, 2, 3, 4, 5]
zip_longest
길이 다른 iterable을 가장 긴 쪽 기준으로.
list(itertools.zip_longest("abc", [1, 2, 3, 4, 5], fillvalue="?"))
# [('a', 1), ('b', 2), ('c', 3), ('?', 4), ('?', 5)]
내장 zip(strict=True)(3.10+)은 길이 불일치 시 ValueError.
누적
accumulate
reduce의 모든 중간 결과를 yield.
import operator
list(itertools.accumulate([1, 2, 3, 4])) # [1, 3, 6, 10]
list(itertools.accumulate([1, 2, 3, 4], operator.mul)) # [1, 2, 6, 24]
list(itertools.accumulate([1, 2, 3, 4], initial=100)) # 3.8+: [100, 101, 103, 106, 110]
그룹핑
groupby
인접한 동일 키 원소를 그룹핑. 정렬되지 않으면 같은 키도 따로 그룹.
python
import itertools
data = "AAABBCDDDD"
for key, group in itertools.groupby(data):
print(key, list(group))
# 키 함수
words = ["apple", "ant", "banana", "blueberry", "cat"]
words.sort(key=lambda s: s[0]) # 같은 첫 글자끼리 인접하게
for key, group in itertools.groupby(words, key=lambda s: s[0]):
print(key, list(group)) 결과
A ['A', 'A', 'A']
B ['B', 'B']
C ['C']
D ['D', 'D', 'D', 'D']
a ['apple', 'ant']
b ['banana', 'blueberry']
c ['cat']중요: 그룹 객체는 iterator라서 다음 그룹으로 넘어가면 사라진다. 보관하려면 list(group) 변환.
조합·순열
product
데카르트 곱.
list(itertools.product([1, 2], "ab"))
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]
list(itertools.product([0, 1], repeat=3))
# [(0,0,0), (0,0,1), (0,1,0), (0,1,1), (1,0,0), (1,0,1), (1,1,0), (1,1,1)]
permutations / combinations / combinations_with_replacement
python
import itertools
print(list(itertools.permutations("ABC", 2)))
print(list(itertools.combinations("ABCD", 2)))
print(list(itertools.combinations_with_replacement("AB", 2))) 결과
[('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]
[('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]
[('A', 'A'), ('A', 'B'), ('B', 'B')]수학 정의 그대로. nPr, nCr, multiset 등.
permutations vs combinations
- permutations: 순서 의미 있음 (
(A, B)≠(B, A)), 중복 불가 - combinations: 순서 무관 (
{A, B}={B, A}), 중복 불가 - combinations_with_replacement: 순서 무관, 같은 원소 여러 번
- product: 순서 의미 있음, 같은 원소 여러 번
분기
tee
하나의 iterator를 N개로 복제. 한쪽이 너무 빠르게 진행하면 내부 버퍼 메모리.
it1, it2 = itertools.tee(range(5), 2)
print(list(it1)) # [0, 1, 2, 3, 4]
print(list(it2)) # [0, 1, 2, 3, 4]
starmap
map은 함수에 원소 하나씩 전달. starmap은 튜플을 unpack해서 전달.
list(map(pow, [(2, 3), (3, 4)])) # TypeError (pow는 2인수)
list(itertools.starmap(pow, [(2, 3), (3, 4)])) # [8, 81]
실용 레시피
itertools 공식 문서의 “Recipes” 섹션에 30+ 패턴이 있다. 자주 쓰이는 몇 가지.
batched (3.12+)
list(itertools.batched("ABCDEFG", 3))
# [('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)]
3.12 이전엔 직접 구현:
def batched(iterable, n):
it = iter(iterable)
while batch := tuple(itertools.islice(it, n)):
yield batch
pairwise (3.10+)
연속된 두 원소 쌍.
list(itertools.pairwise([1, 2, 3, 4]))
# [(1, 2), (2, 3), (3, 4)]
# 차이 계산
diffs = [b - a for a, b in itertools.pairwise(nums)]
중복 제거 유지 순서
def unique_everseen(iterable, key=None):
seen = set()
for x in iterable:
k = key(x) if key else x
if k not in seen:
seen.add(k)
yield x
list(unique_everseen("AAAABBBCCDAABBB")) # ['A', 'B', 'C', 'D']
함정
1. iterator 1회용
g = itertools.chain([1, 2], [3, 4])
list(g) # [1, 2, 3, 4]
list(g) # [] (소진됨)
여러 번 순회하려면 list로 변환 후 사용 또는 매번 새로 생성.
2. tee의 메모리
it1, it2 = itertools.tee(big_iter)
process_all(it1) # 끝까지 진행
do_other(it2) # 처음부터 다시
# it2가 진행하는 동안 it1이 만든 원소가 메모리에 쌓임
두 사용자가 비슷한 속도로 진행해야 효율적. 아니면 list 변환이 나음.
3. groupby 정렬 필수
data = "ABABAB"
list(itertools.groupby(data)) # 그룹 6개 (인접 동일만 묶음)
전체 동일 키 그룹핑은 sorted 후 groupby.
💬 댓글