[python/알고리즘] itertools 라이브러리 combinations로 모든 조합 찾기

itertools 라이브러리

combinations

• combinations(iterable, n): iterable에 들어있는 값 중 n개를 뽑을 수 있는 조합 반환

from itertools import combinations

alist = ['a', 'b', 'c']

for i in combinations(alist, 2):
	print(i)

# 결과:
('a', 'b')
('a', 'c')
('b', 'c')

combinations_with_replacement

• combinations_with_replacement(iterable, n): 복원 추출이 가능할 때 조합 반환

from itertools import *

alist = ['a', 'b', 'c']

for i in combinations_with_replacement(alist, 2):
	print(i)
    
# 결과:
('a', 'a')
('a', 'b')
('a', 'c')
('b', 'b')
('b', 'c')
('c', 'c')

permutations

• permutations(iterable, n): 원소 개수가 n인 순열 반환

from itertools import *

alist = ['a', 'b', 'c']

for i in permutations(alist, 2):
	print(i)
    
# 결과:
('a', 'b')
('a', 'c')
('b', 'a')
('b', 'c')
('c', 'a')
('c', 'b')

itertools 활용하여 문제 풀기

프로그래머스 코드카타 42번 - 삼총사

itertools 없이 짠 코드

from math import *

def solution(number):
    answer = 0
    for i in range(len(number)):
        for j in range(i + 1, len(number)):
            for k in range(j + 1, len(number)):
                if number[i] + number[j] + number[k] == 0:
                    answer += 1
                
    return answer

3중 for문을 돌면서 모든 조합 확인

itertools를 이용하여 짠 코드

from itertools import combinations

def solution(number):
    answer = 0
    for i in combinations(number, 3):
        if sum(i) == 0:
            answer += 1
    return answer

훨씬 짧고 간결해지는 것을 알 수 있다.

참고자료 - 파이썬 공식문서 itertools