Python 내장 패키지 itertools 도장깨기

코딩 테스트를 볼 때 내장 패키지를 알면 도움이 되는 경우가 적지 않다.

외부 패키지를 import할 수 없기 때문에, 내장 패키지를 최대한 활용하는 것이 좋다.

특히 Python의 경우 언어가 단순한 만큼 input이 커지면 시간복잡도가 다른 언어에 비해 높다.

따라서 직접 짠 모듈을 쓰는 것보다 기존의 내장 패키지를 import하는 것이 좋다.

오늘은 그 중 itertools라는 iterable한 객체를 처리할 수 있는 패키지에 대해 정리해보려 한다.

 

공식 package guideline은 아래와 같다.

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

itertools — 효율적인 루핑을 위한 이터레이터를 만드는 함수 — Python 3.9.6 문서

차근 차근 살펴보자.

 

 

무한 이터레이터 - 무한히 반복되는 경우

1. count

count 함수는 말 그대로 숫자를 세는 이터레이터이다.

해당 숫자로부터 무한히 새로운 숫자를 생성해낸다.

음의 방향으로 세는 것도 가능하다.

import itertools

# 10부터 1씩 100까지 세기
for i in itertools.count(10):
    print(i)
    if i >= 100:
        break

# 10부터 5씩 100까지 세기
for i in itertools.count(10,5):
	print(i)
    if i >= 100:
    	break

# 10부터 -2씩 0까지 세기
for i in itertools.count(10,-2):
    print(i)
    if i == 0:
        break

 

2. cycle

객체를 무한히 반복하는 이터레이터이다.

아래의 예시를 보면 이해가 쉽다.

import itertools

# ABCD를 각각 100번 반복하기
n = 0
for i in itertools.cycle('ABCD'):
    n += 1
    print(i)
    
    if n >= 100:
        break

A
B
C
D
A
B
C
D

 

3. repeat

같은 element를 무한히 반복하는 이터레이터이다.

import itertools

# ABCD 100번 반복하여 출력하기
n = 0
for i in itertools.repeat('ABCD'):
    n += 1
    print(i)
    
    if n >= 100:
        break

ABCD
ABCD
ABCD
ABCD
ABCD
ABCD
ABCD
ABCD
ABCD

 

 

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

용도가 가장 많은 경우의 이터레이터들이다.

1. accumulate

누적 합을 계산해주는 이터레이터이다.

import itertools

list(itertools.accumulate([1,2,3,4,5]))

[1, 3, 6, 10, 15]

 

2. chain

객체를 받아 객체들을 나열해주는 함수이다.

특히 해당 함수는 2중 리스트를 flatten할 때 많이 사용한다.

import itertools

list(itertools.chain('ABC','DEF'))

['A', 'B', 'C', 'D', 'E', 'F']

 

import itertools

# 이중 리스트 flatten하기
list_double = [[1,2,3,4],['A','B','C','D']]
list(itertools.chain(*list_double))

[1, 2, 3, 4, 'A', 'B', 'C', 'D']

 

3. chain.from_iterable

chain과 같은 역할을 하지만 iterable한 객체를 input으로 받는다. (iterable)

import itertools

# iterable한 객체인 list를 인풋으로 받음
list(itertools.chain.from_iterable(['ABC', 'DEF']))

['A', 'B', 'C', 'D', 'E', 'F']

 

4. compress

요소를 필터링하여 selector에서 True로 평가되는 요소만 반환함 (data, selectors)

import itertools

# True라고 판명되는 인자들만 반환
list(itertools.compress('ABCDEF', [1,0,1,0,1,1]))

['A', 'C', 'E', 'F']

 

5. dropwhile

pred가 실패할 시 시작 (func, seq)

import itertools

# 실패 시 시작되는 요소들 반환
list(itertools.dropwhile(lambda x: x<5, [1,4,6,4,1]))

[6, 4, 1]

 

6. filterfalse

거짓인 요소들 반환 (func, seq)

import itertools

list(itertools.filterfalse(lambda x: x%2, range(10)))

[0, 2, 4, 6, 8]

 

7. groupby

key값을 그룹화하여 iterable하게 출력해주는 함수.

key name과 grouper iterable object를 반환한다.

d = [('Europe','Manchester'),
     ('America','NewYork'),
     ('Asia','Seoul'),
     ('Asia','Tokyo'),
     ('America','Chicago'),
     ('America','Seattle'),
     ('Europe','London'),
     ('Asia','Beijing'),
     ('Europe','Paris'),
     ]

category = {}
for k, g in itertools.groupby(sorted(d), lambda x:x[0]):
    print(k)
    print(list(g))

America
[('America', 'Chicago'), ('America', 'NewYork'), ('America', 'Seattle')]


Asia
[('Asia', 'Beijing'), ('Asia', 'Seoul'), ('Asia', 'Tokyo')]


Europe
[('Europe', 'London'), ('Europe', 'Manchester'), ('Europe', 'Paris')]

 

8. islice

iteration에 있어 시작점을 정하여 시작 가능한 함수.

for문에서 시작점을 정해줄 때 많이 사용한다.

range함수의 경우 range(2,10)과 같이 시작점을 정해줄 수 있지만,

일반적인 iterable 객체의 경우 그럴 수 없는 경우가 많다.

이 때 사용하면 굉장히 유용하다.

import itertools

# iterable 객체를 3번째 요소부터 반복시킴
for i in itertools.islice('ABCDEFG', 2, None):
    print(i)

C
D
E
F
G

 

9. startmap

function을 mapping하여 적용시켜주는 함수

import itertools

# (x,y) x^y함수를 매핑하여 list로 반환하기
list(itertools.starmap(pow, [(2,5), (3,2), (10,3)]))

[32, 9, 1000]

 

10. takewhile

dropwhile과 반대되는 함수.

시작할 때부터 False가 나오는 지점 (실패할 때까지) 반환

import itertools

# 5이상인 숫자가 나오면 반환 중지
list(itertools.takewhile(lambda x: x<5, [1,4,6,4,1]))

[1, 4]

 

11. tee

이터레이터를 n개로 복사

import itertools

iter_obj = range(5,10)

list1, list2, list3 = itertools.tee(iter_obj, 3)

print(list(list1))
print(list(list2))
print(list(list3))

[5, 6, 7, 8, 9]
[5, 6, 7, 8, 9]
[5, 6, 7, 8, 9]

 

12. zip_longest

iterable한 객체를 각각 paste함. (p, q, fillvalue)

import itertools

list(itertools.zip_longest('ABCD', 'xy', fillvalue='-'))

[('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')]

 

 

조합형 이터레이터

조합, 순열 등의 기능을 제공하는 함수이다.

1. product

데카르트 곱. 모든 객체를 하나씩 조합한다.

import itertools

list(itertools.product('ABCD', repeat=2))

[('A', 'A'),
 ('A', 'B'),
 ('A', 'C'),
 ('A', 'D'),
 ('B', 'A'),
 ('B', 'B'),
 ('B', 'C'),
 ('B', 'D'),
 ('C', 'A'),
 ('C', 'B'),
 ('C', 'C'),
 ('C', 'D'),
 ('D', 'A'),
 ('D', 'B'),
 ('D', 'C'),
 ('D', 'D')]

 

2. permutation

순열 결과를 반환한다.

import itertools

list(itertools.permutations('ABCD', 2))

[('A', 'B'),
 ('A', 'C'),
 ('A', 'D'),
 ('B', 'A'),
 ('B', 'C'),
 ('B', 'D'),
 ('C', 'A'),
 ('C', 'B'),
 ('C', 'D'),
 ('D', 'A'),
 ('D', 'B'),
 ('D', 'C')]

 

3. combination

조합 결과를 반환한다.

import itertools

list(itertools.combinations('ABCD', 2))

[('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]

 

4. combinations_with_replacement

복원 추출 결과를 반환한다.

import itertools

list(itertools.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')]