본문 바로가기

파이썬

[Python3] 효율적인 반복자 itertools 사용하기

import itertools

효율적인 루핑을 위한 이터레이터를 만드는 함수

파이썬의 공식 문서는 itertools 모듈을 위와 같이 소개한다. 한 마디로 자신만의 반복자를 만드는 모듈이다.

https://docs.python.org/ko/3.8/library/itertools.html

아래는 위의 문서를 읽고 직접 코드를 실행하며 정리한 내용이다.

무한 이터레이터

itertools.count(start, [step])

>>> for i in itertools.count(10,2):
...     print(i)
...     if(i==20):break
...
10
12
14
16
18
20

시작 값(start)와 step마다 증가할 값(step)을 인자로 하여 무한하게 숫자를 생성한다.

break를 걸지 않으면 무한 루프를 돌게 되므로 주의한다.

 

itertools.cycle(p)

>>> for idx, word in enumerate(itertools.cycle('abcd')):
...     print(word)
...     if idx==10 : break
...
a
b
c
d
a
b
c
d
a
b
c

p에는 문자열, 리스트, 튜플, 딕셔너리 모두 들어갈 수 있다.

인덱스 순서대로 모든 인자를 순서대로 무한 출력한다.

itertools.repeat(elem, [n])

>>> for idx, word in enumerate(itertools.repeat({1:2},3)):
...     print(word)
...
{1: 2}
{1: 2}
{1: 2}

반복될 횟수를 정해주지 않으면 무한으로 반복한다. 지정된 요소만 무한으로 반복한다.

짧은 입력 시퀀스에서 종료되는 이터레이터

itertools.accumulate(p)

>>> for i in itertools.accumulate([1,2,3,4]):
...  print(i)
...
1
3
6
10

누적 결과합을 반환한다. 1, 1+2, 1+2+3, 1+2+3+4.

itertools.chain(p,q,...)

>>> for i in itertools.chain([1,2,3,4],(5,6,7)):
...     print(i)
...
1
2
3
4
5
6
7

서로 다른 타입이라도 각 리스트를 연결해준다.

itertools.chain.from_iterable()

>>> for i in itertools.chain.from_iterable(['abc','def']):
...     print(i)
...
a
b
c
d
e
f

리스트의 요소끼리 연결해서 하나씩 리턴해준다.

itertools.compress(data, selectors)

>>> for i in itertools.compress('himyname',[1,0,0,1,0,1,1]):
...     print(i)
...
h
y
a
m

데이터에서 원하는 인덱스만 뽑아낼 수 있다.

itertools.dropwhile(pred, seq)

>>> for i in itertools.dropwhile(lambda x: x<5, [2,3,4,6,1,2]):
...     print(i)
...
6
1
2

 

lambda 식의 조건이 만족되지 않는 (거짓이 되는) 순간부터 시작된다.

itertools.takewhile()

>>> for i in itertools.takewhile(lambda x: x<5, [1,4,2,3,5,6,1,2,4]):
...     print(i)
...
1
4
2
3

dropwhile()과 반대로 조건이 만족되지 않는 순간까지 요소를 리턴한다.

itertools.groupby()

>>> from operator import itemgetter

>>> t = [('apple', 2), ('apple', 5), ('banana',12), ('orange', 53)]
>>> for key, value in itertools.groupby(t, key=itemgetter(0)):
...     print(key, list(map(itemgetter(1), value)))
...
apple [2, 5]
banana [12]
orange [53]

key함수를 이용하여 첫번째 요소로 groupby 한다. 이때 자료는 정렬되어 있어야 한다.

만약 정렬되어 있지 않을 경우, t.sort(key=itemgetter(0)) 으로 정렬가능하다.

itertools.tee(iter,n)

>>> for i in itertools.tee('abdfs',2):
...     for j in i:
...             print(j)
...
a
b
d
f
s
a
b
d
f
s

하나의 iterator을 n개의 iterator로 나눈다.

itertools.starmap(func, seq)

>>> for i in itertools.starmap(pow, [(2,5), (3,2), (10,10)]):
...     print(i)
...
32
9
10000000000

함수를 적용하여 리턴하는 반복자로, map, filter 함수와 비슷한 역할이다.

 

조합형 이터레이터 (가장 자주 쓰임)

itertools.product(p,q, [repeat=1])

>>> for i in itertools.product([1,2,3],[4,5,6]):
...     print(i)
...
(1, 4)
(1, 5)
(1, 6)
(2, 4)
(2, 5)
(2, 6)
(3, 4)
(3, 5)
(3, 6)
>>> for i in itertools.product([1,2],[4,5], repeat=2):
...     print(i)
...
(1, 4, 1, 4)
(1, 4, 1, 5)
(1, 4, 2, 4)
(1, 4, 2, 5)
(1, 5, 1, 4)
(1, 5, 1, 5)
(1, 5, 2, 4)
(1, 5, 2, 5)
(2, 4, 1, 4)
(2, 4, 1, 5)
(2, 4, 2, 4)
(2, 4, 2, 5)
(2, 5, 1, 4)
(2, 5, 1, 5)
(2, 5, 2, 4)
(2, 5, 2, 5)

데카르트 곱으로 중첩된 for 루프와 동일한 기능을 한다.

만약 문자열 'ABC' 였다면, AA, AB, AC, BA, BB, BC, CA, CB, CC 의 결과가 반환된다.

itertools.permutations(p,r)

>>> for i in itertools.permutations([1,2,3],2):
...     print(i)
...
(1, 2)
(1, 3)
(2, 1)
(2, 3)
(3, 1)
(3, 2)

nPr 의 순열을 표현할 수 있는 이터레이터이다.

만약 문자열 'ABC' 였다면, AB, AC, BA, BC, CA, CB 의 결과가 반환된다.

itertools.combinations(p,r)

>>> for i in itertools.combinations([1,2,3],2):
...     print(i)
...
(1, 2)
(1, 3)
(2, 3)

nCr 의 조합을 표현할 수 있는 이터레이터이다. (반복되는 요소 없음)

만약 문자열 'ABC' 였다면, AB, AC, BC의 결과가 반환된다.

itertools.combinations_with_replacement(p,r)

>>> for i in itertools.combinations_with_replacement([1,2,3],2):
...     print(i)
...
(1, 1)
(1, 2)
(1, 3)
(2, 2)
(2, 3)
(3, 3)

nCr의 조합을 표현할 수 있는 이터레이터이다.(반복 되는 요소 있음)

만약 문자열 'ABC' 였다면, AA, AB, AC, BB, BC, CC 의 결과가 반환된다.