본문으로 건너뛰기
김신건의 로그

[Python] itertools: 효율적인 iterator 빌딩 블록

· 수정 · 📖 약 1분 · 516자/단어 #python #itertools #stdlib #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개 (인접 동일만 묶음)

전체 동일 키 그룹핑은 sortedgroupby.

이 개념을 다룬 위키 페이지 (1)

💬 댓글

사이트 검색 / 명령어

검색

스크롤 = 확대/축소 · 드래그 = 이동 · 0 = 원래 크기 · ESC = 닫기