파이썬 itertools 모듈

코딩을 하다보면 array, tuble, dictionary 등의 이터러블을 다뤄야 하거나 동일한 동작을 반복적으로 수행해야하는 경우가 많은데, 이 경우 유용하게 이터레이터를 생성할 수 있는 게 파이썬의 itertools 모듈이다.

 

 

0. 이터러블과 이터레이터

이터러블(iterable)이란 반복 가능한 객체를 의미하며, 이터레이터(iterator)란 이터러블에서 값을 생성해낼 수 있는 객체를 의미한다.

 

파이썬 내장함수 중 iter()를 사용해 이터러블을 인자로 하여 이터레이터 객체를 만들어 낼 수 있고, next()를 이용해 이터레이터를 호출하면 이터러블의 요소를 차례대로 꺼낼 수 있다.

어떤 타입의 객체가 이터러블인지 확인하는 가장 쉬운 방법은 iter()의 인자로 해당 객체를 전달하는 것이다. 이터러블이 아니면 에러가 난다.

아래의 코드에서 x는 이터러블에 해당하고, y는 이터러블로부터 값을 생성해내는 이터레이터의 인스턴스이다.

x = 'ab'
y = iter(x)

print(list(i for i  in list(y)))
>>> ['a', 'b']

print(next(y))
>>> a
print(next(y))
>>> b
print(next(y))
>>> StopIteration

 

 

1. 무한 이터레이터

count(start, [step]) : start 부터 시작해서 step씩 증가한 요소를 반복적으로 생성한다. step을 지정하지 않으면 start부터 1씩 증가한다.

from itertools import count

for i in count(1, 0.5):
    if i < 3:
    	print(i, end = '  ')
    else:
    	break

>>> 1  1.5  2.0  2.5

 

cycle(p) : 이터러블 객체 p에 대해 p의 각 요소를 반복적으로 생성해준다.

from itertools import cycle

count = 1
for i in cycle(('a', 'c', 'b')):
    print(i, end = '  ')
    count += 1
    if count > 5: 
        break

>>> a  c  b  a  c

 

repeat(elem, [n]) : elem을 n번 반복하여 생성해준다. n을 지정하지 않을 경우 무한으로 생성해준다.

from itertools import repeat

for i in repeat('test', 3):
	print(i, end = '  ')
    
>>> test  test  test

 

 

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

accumulate(p, [func], [initial=None]) : 이터러블 객체 p에 대해 각 요소까지의 누적합을 반환해준다.

 

선택적 func 인자를 통해 다른 이항함수를 지정해주면 각 요소까지의 누적 결과를 반환해준다. 예를 들어, 두 번째 인자로 max를 입력해주면 해당 요소까지의 값들 중 최댓값을 구할 수 있다.

 

또한, 키워드 인자를 통해 initial값을 지정할 수 있는데 이 경우 p의 첫 번째 요소로 initial값이 추가되는 것과 같은 결과가 나온다. 최종적으로 출력되는 갯수 역시 입력된 이터러블의 요소 수 +1이 된다.

from itertools import accumulate
	
test = [1, -2, 3, -4, 5]

print(*accumulate(test))
>>> 1 -1 2 -2 3

print(*accumulate(test, max))
>>> 1 1 3 3 5

print(accumulate(test, min))
>>> 1 -2 -2 -4 -4

print(accumulate(test, initial=100))
>>> 100 101 99 102 98 103

 

chain(p, q, ...) : 여러 이터러블 객체들을 연결해준다.

from itertools import chain

print(*chain('ABC', 'DE'))
>>> A B C D E

print(*chain(['A', 'B', 'C'], ['D', 'E']))
>>> A B C D E

 

compress(data, selectors) : 동일 인덱스에 대해 selectors의 값이 1인 data 요소만 반환해준다.

from itertools import compress

print(*compress(['a', 'b', 'c'],[1, 0, 1]))
>>> a c

 

takewhile(pred, seq) : pred가 거짓일 때까지 seq 요소들을 반환해준다.

dropwhile(pred, seq) : pred가 거짓일 때부터 seq 요소들을 반환해준다.

filterfalse(pred, seq) : pred가 거짓인 seq 요소들을 반환해준다.

from itertools import takewhile, dropwhile, filterwhile

test = [1, 2, 3, 2, 1]

print(*takewhile(lambda x: x<3, test))
>>> 1 2

print(*dropwhile(lambda x: x<3, test))
>>> 3 2 1

print(*filterfalse(lambda x: x<3, test))
>>> 3

 

 

 

3. 조합형 이터레이터

이터레이터	순서	중복	예시(인자 -> 결과)
product()	O	O	'ABC', repeat=2	AA AB AC BA BB BC CA CB CC
permutations()	O	X	'ABC', 2	AB AC BA BC CA CB
combinations()	X	X	'ABC', 2	AB AC BC
combinations_with_replacement()	X	O	'ABC', 2	AA AB AC BB BC CC

사실 이 조합형 이터레이터들에 대해 정리하고 싶어서 이 글을 쓰기 시작했던 건데, 찾아보니 itertools 모듈의 다른 함수들이 생각보다 훨씬 유용할 것 같아서 다른 이터레이터들에 대해서도 정리하게 되었다.

 

경우의 수를 따질 때 대표적으로 쓰이는 조합론적인 구조로 순열과 조합이 있다.

순열이란 서로 다른 n개의 원소에서 r개를 중복없이 순서에 상관 있게 선택하는 것을 의미하며, 조합이란 서로 다른 n개의 원소에서 r개를 중복없이 순서를 고려하지 않고 선택하는 것을 의미한다.

 

파이썬에서는 permutations()를 통해 주어진 이터러블에 대한 모든 순열을 구할 수 있고, combinations()를 통해 주어진 이터러블에 대한 모든 조합을 구할 수 있다.

 

또한, product()를 통해 주어진 모든 이터러블에 대한 모든 순서쌍(데카르트곱)을 구할 수 있으며, combinations_with_replacement()를 통해 주어진 이터러블에 대해 중복을 허용한 모든 조합을 구할 수도 이다.

 

 

product(p, q, ..., [repeat=1]) : 주어진 이터러블 p, q, ...에 대하여 순서와 중복을 고려한 모든 조합을 반환해준다. 키워드인자로 반복횟수 repeat지정해주면 해당 횟수만큼 전체 이터러블이 반복되어 조합된다.

from itertools import product

test1 = [1, 2, 3]
test2 = [4, 5]

print(*product(test1, repeat=2))
>>> (1, 1) (1, 2) (1, 3) (2, 1) (2, 2) (2, 3) (3, 1) (3, 2) (3, 3)

print(*product(test1, test2))
>>> (1, 4) (1, 5) (2, 4) (2, 5) (3, 4) (3, 5)

 

permutations(p, [r]) : 주어진 이터러블 p에 대하여 중복 없이 r개의 원소로 만들 수 있는 모든 순열을 반환해준다. r을 별도로 지정해주지 않을 경우 전체 p에 대한 가능한 모든 순열을 반환해준다.(즉, r=len(p))

from itertools import permutaions

test1 = [1, 2, 3]

print(*permutations(test1))
>>> (1, 2, 3) (1, 3, 2) (2, 1, 3) (2, 3, 1) (3, 1, 2) (3, 2, 1)

print(*permutations(test1, 2))
>>> (1, 2) (1, 3) (2, 1) (2, 3) (3, 1) (3, 2)

 

combinations(p, r) : 주어진 이터러블 p에 대하여 중복 없이 r개의 원소로 만들 수 있는 모든 조합을 반환해준다.

from itertools import combinations

test1 = [1, 2, 3]

print(*combinations(test1, 2))
>>> (1, 2) (1, 3) (2, 3)

 

combinations_with_replacement(p, r) : 주어진 이터러블 p에 대하여 중복을 허용하여 r개의 원소로 만들 수 있는 모든 조합을 반환해준다.

from itertools import combinations_with_replacement

test1 = [1, 2, 3]

print(*combinations_with_replacement(test1, 2))
>>> (1, 1) (1, 2) (1, 3) (2, 2) (2, 3) (3, 3)

 

 

참고사이트:

itertools --- 효율적인 루핑을 위한 이터레이터를 만드는 함수 — 파이썬 설명서 주석판 (flowdas.com)

itertools --- 효율적인 루핑을 위한 이터레이터를 만드는 함수 — 파이썬 설명서 주석판