itertools 라이브러리란?

itertools
단순한 루핑을 효율적으로 만들어주는 라이브러리

DAY21 의 TIL 에서 itertools 의 용도를 간략하게 설명했었으나 정확한 사용법에 대해서는 찾아보지 못했었다. 이번에는 해당 라이브러리에 들어있는 함수들을 좀 더 살펴보기로 했다.
DAY21 에서 언급한 'from itertools import product' 는 itertools 의 함수 중 하나에 불과했다.

---

함수 종류

combinations

입력 iterable에서 요소를 중복 없이 r개씩 뽑아 가능한 모든 조합을 생성한다.
순서를 고려하지 않기에 (m,n) 과 (n,m) 을 중복된 데이터로 판단한다.

from itertools import combinations

data = [1, 2, 3]
result = list(combinations(data, 2))
print(result)

#결과: [(1, 2), (1, 3), (2, 3)]

permutations

입력 iterable에서 요소를 중복 없이 r개씩 뽑아 가능한 모든 순열을 생성한다.
순서를 고려하기에 (m,n) 과 (n,m) 을 다른 데이터로 판단한다.

from itertools import permutations

data = [1, 2, 3]
result = list(permutations(data, 2))
print(result)

#결과: [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]

product

주어진 iterable들로 가능한 모든 곱을 도출한다.
repeat는 각 iterable을 얼마나 반복할지 결정한다.

from itertools import product

a = [1, 2]
b = [3, 4]
result = list(product(a, b))
print(result)

#결과: [(1, 3), (1, 4), (2, 3), (2, 4)]

combinations_with_replacement

combinations_with_replacement(iterable, r) 과 같은 구조를 쓴다.
iterable에서 요소를 중복을 허용하여 r개씩 뽑아 가능한 모든 조합을 도출한다.

from itertools import combinations_with_replacement

data = [1, 2, 3]
result = list(combinations_with_replacement(data, 2))
print(result)

#결과: [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

chain

여러 iterable을 인수로 받아 하나의 iterable처럼 동작하는 반복자를 생성한다.
iterable들을 합치기 위해 쓸 수 있다.

from itertools import chain

data1 = [1, 2, 3]
data2 = ['a', 'b', 'c']
result = list(chain(data1, data2))
print(result)

#결과: [1, 2, 3, 'a', 'b', 'c']

cycle

입력 iterable을 무한히 반복하는 반복자를 생성한다.

from itertools import cycle

data = [1, 2, 3]
inf_cycle = cycle(data)
for i in range(10):
    print(next(inf_cycle))
    
#결과:
1
2
3
1
2
3
1
2
3
1

accumulate

iterable의 누적 합을 계산하여 반환한다. func 인수를 통해 다른 누적 함수를 사용할 수 있다.

from itertools import accumulate
import operator

data = [1, 2, 3, 4]

# 앞에서부터 순서대로 덧셈 (1=1, 1+2=3, 3+3=6, 6+4=10)

result = list(accumulate(data))
print(result)

# operator.mul 붙일 시 곱셈으로 변환 (1=1, 1*2=2, 2*3=6, 6*4=24)

result = list(accumulate(data, operator.mul))
print(result)

groupby

iterable을 key 함수로 그룹화한다.

from itertools import groupby

data = sorted([('a', 1), ('b', 2), ('a', 3), ('b', 4)], key=lambda x: x[0])
grouped = groupby(data, key=lambda x: x[0])
for key, group in grouped:
    print(key, list(group))
    
#결과:
a [('a', 1), ('a', 3)]
b [('b', 2), ('b', 4)]

---

문제풀이

[문제]
한국중학교에 다니는 학생들은 각자 정수 번호를 갖고 있습니다. 이 학교 학생 3명의 정수 번호를 더했을 때 0이 되면 3명의 학생은 삼총사라고 합니다. 예를 들어, 5명의 학생이 있고, 각각의 정수 번호가 순서대로 -2, 3, 0, 2, -5일 때, 첫 번째, 세 번째, 네 번째 학생의 정수 번호를 더하면 0이므로 세 학생은 삼총사입니다. 또한, 두 번째, 네 번째, 다섯 번째 학생의 정수 번호를 더해도 0이므로 세 학생도 삼총사입니다. 따라서 이 경우 한국중학교에서는 두 가지 방법으로 삼총사를 만들 수 있습니다.
한국중학교 학생들의 번호를 나타내는 정수 배열 number가 매개변수로 주어질 때, 학생들 중 삼총사를 만들 수 있는 방법의 수를 return 하도록 solution 함수를 완성하세요.

[제한사항]

• 3 ≤ number의 길이 ≤ 13
• -1,000 ≤ number의 각 원소 ≤ 1,000
• 서로 다른 학생의 정수 번호가 같을 수 있습니다.

[풀이]

from itertools import combinations
def solution(number):
    return sum(1 for combo in combinations(number, 3) if sum(combo) == 0)

1. combinations 를 이용해 입력받은 number 로부터 3가지 요소를 꺼내는 것을 반복한다.
2. 3가지 요소는 combo 안에 들어간다.
3. combo 의 합이 0인 경우 return 값에 1이 더해진다.
4. combinations 와 for 구문의 조합에 의해 numbers 에서 3가지 요소를 꺼내 만들 수 있는 모든 조합이 도출되면 for 구문의 작동은 중단한다.
5. 이 시점에서 return 값은 이 조합들 중 '셋이 합쳐 0이 되는 경우' * 1 이 된다.