[Python] itertools 프로그래밍

파이썬 내장 라이브러리인 itertools는 특정배열에 대하여 순열이나 조합을 이용한 문제를 풀 때 유용하며 자신의 반복자를 만드는 모듈이다. 이 라이브러리에는 상당히 많은 함수들이 내장되어 있기에, 알아두고 가는 것이 많이 도움이 된다.

 

1. cycle

반복 가능한 객체를 순서대로 무한히 반복하는 iterator 생성

import itertools

student = itertools.cycle(['가나다','라마바','사아자'])

next(student)

# 가나다
# 라마바
# 사아자
# 가나다
# 라마바
...

 

2. accumulate

반복 가능한 객체의 누적 합을 계산하여 iterator로 반환하는 함수

보통은 iterator를 반환하기 때문에 아래와 같이 list로 깔끔하게 반환한다.

monthly_income = [250, 259, 260, 310, 234, 410, 371, 381, 295, 276, 285, 398]
result = list(itertools.accumulate(monthly_income))

print(result)
# [250, 509, 769, 1079, 1313, 1723, 2094, 2475, 2770, 3046, 3331, 3729]

 

3. groupby

반복 가능한 객체를 키값으로 분류하고, 그 결과를 반환하는 함수

groupby 함수를 사용하기 전에 먼저 데이터(객체)는 정렬되어 있어야 한다.

그렇지 않으면, 아래와 같이 클래스(key)가 하나로 모이지 않고 뿔뿔이 흩어진 채 나타나게 된다.(목적 실패)

data = [
	{'name':'alice', 'class':'6'},
    {'name':'ornn', 'class':'4'},
    {'name':'aatrox', 'class':'3'},
    {'name':'jax', 'class':'6'},
    {'name':'seraphin', 'class':'1'},
    {'name':'kindred', 'class':'2'},
    {'name':'ari', 'class':'4'}
]

data = sorted(data, key=operator.itemgetter('class'))
grouped_data = itertools.groupby(data,key=operator.itemgetter('class'))

result = {}
for key, group_data in grouped_data:
    result[key] = list(group_data)
    
pprint.pprint(result)

#{'1': [{'class': '1', 'name': 'seraphin'}],
# '2': [{'class': '2', 'name': 'kindred'}],
# '3': [{'class': '3', 'name': 'aatrox'}],
# '4': [{'class': '4', 'name': 'ornn'}, {'class': '4', 'name': 'ari'}],
# '6': [{'class': '6', 'name': 'alice'}, {'class': '6', 'name': 'jax'}]}

 

4. permutation

반복 가능 객체 중에서 n개를 선택한 순열을 반환하는 함수

순열과 조합에 대해 구분을 잘 해야 한다.

from itertools import permutations

data = [1, 2, 3]
permutations_result = permutations(data)

for perm in permutations_result:
    print(perm)

(1, 2, 3)
(1, 3, 2)
(2, 1, 3)
(2, 3, 1)
(3, 1, 2)
(3, 2, 1)

 

5. cominations

반복 가능 객체 중에서 n개를 선택한 조합을 반환하는 함수

from itertools import combinations

data = [1, 2, 3]
combinations_result = combinations(data, 2)

for combo in combinations_result:
    print(combo)
    
    
(1, 2)
(1, 3)
(2, 3)

 

6. count

반복해야할 최대 수까지의 범위를 정하지 않았을 때 사용되며 count(a,b)인 경우 a에서 시작해서 b씩 증가하는 순으로 반복

from itertools import count

for i in count(1,5):
    if i <= 21:
        print(i)
    else:
        break

for Num, Chr in zip(count(0, 10), ['a', 'b', 'c', 'd', 'e']):
    """
    zip()은 동일한 개수로 이루어진 자료형을 묶어주는 역할
    """
    print('{0}: {1}'.format(Num, Chr))
    
    
1
6
11
16
21
0: a
10: b
20: c
30: d
40: e

 

이외에도 많지만, 일단 코테에서 많이 본 것들은 다음과 같아서 정리했다.