Itertools 순열 조합 라이브러리

알고리즘 문제를 풀다보면 순열이나 조합을 사용해서 문제를 풀 때 더 간단하게 풀 수 있는 경우가 많다. 그럴 때에는 itertools 라이브러리를 사용한다.
itertools 라이브러리에는 사실 순열과 조합 외에도 iterable 자료형을 처리하는 다른 함수들이 존재하지만 가장 자주 쓰이는 것이 순열과 조합이기 때문에 여기서는 순열, 조합을 사용하는 법만 확인한다.
다른 기능을 알고 싶으면 https://docs.python.org/ko/3/library/itertools.html 를 확인하길 바란다.

1. 순열 Permutation

itertools.permutations(iterable, r = None)
순열은 iterable 자료형에서 r 개의 데이터를 뽑아 순서가 있는 형태로 나타낸다.
각각의 sub sequence는 길이가 r 이고 튜플로 이루어진다. 반환형태는 클래스이기 때문에 사용하기 힘들 수 있기 때문에 주로 리스트로 변환을 해주는 것이 편하다. (필수는 아니다.)

from itertools import permutations

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

result = list(permutations(data))
result2 = permutations(data)
# 두 번째 인자를 넣어주지 않으면 자동으로 r = len(data) 가 된다.
print(result)
print(result2)
for i in result2:
    print(i, end=' ')

| [('a', 'b', 'c'), ('a', 'c', 'b'), ('b', 'a', 'c'), ('b', 'c', 'a'), ('c', 'a', 'b'), ('c', 'b', 'a')]
<itertools.permutations object at 0x0000020AD7E80EA0>
('a', 'b', 'c') ('a', 'c', 'b') ('b', 'a', 'c') ('b', 'c', 'a') ('c', 'a', 'b') ('c', 'b', 'a') |
| --- |

코드에서 처럼 for 문에서 튜플을 하나씩 가져와서 사용하는 것은 가능하지만 result2[0]같이 직접적으로 객체에 접근하는 것은 불가능하다. 그래서 list로 변환을 해주는 것이 편하다는 것이고 앞으로 나올 combinations, product, combinations_with_replacement 에서는 설명을 생략한다.

2. 조합 Combination

itertools.combinations(iterable, r =None)
조합은 순열은 iterable 자료형에서 r 개의 데이터를 뽑아 순서가 없는 상태로 나열한다.

from itertools import combinations

data = ['a', 'b', 'c']
result = list(combinations(data, 2))    # r = 2 로 설정

print(result)

3. 데카르트 곱(cartesian product)

itertools.product(iterable, repeat = 1)
데카르트 곱은 iterable 자료형에서 repeat 개 만큼을 중복을 허용해서 순열을 만든다고 책에 쓰여져 있기도 하지만 사실 product는 여러개의 iterable 자료형을 입력으로 받아서 가능한 모든 곱을 나타내는 것이다. 다음은 3가지 예시를 보여준다.

from itertools import product

data = ['a', 'b', 'c']
data2 = ['x', 'y']
data3 = range(2)    # [0, 1]
result = list(product(data, repeat=2))

print(result)
print(list(product(data, data2)))
print(list(product(data3, repeat=3)))

| [('a', 'a'), ('a', 'b'), ('a', 'c'), ('b', 'a'), ('b', 'b'), ('b', 'c'), ('c', 'a'), ('c', 'b'), ('c', 'c')]
[('a', 'x'), ('a', 'y'), ('b', 'x'), ('b', 'y'), ('c', 'x'), ('c', 'y')]
[(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)] |
| --- |
| 'product(data, repeat=2) 는 사실 product(data, data) 이다. 두 번째 예시에서 보듯이 data에서의 값과 data2에서의 값으로 만들 수 있는 모든 순열을 만들어 내는 것이다. |
| 세 번째 예시는 나중에 이진수를 만들어 낸다면 유용하게 쓰일 것 같아서 적어봤다. |

4. 중복 조합

** itertools.combinations_with_replacement(iterable, r = None)**
중복 조합은 말 그대로 중복을 허용해서 조합을 완성시키는 것이다. 바로 예시로 넘어간다.

from itertools import combinations_with_replacement

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

result = list(combinations_with_replacement(data, 2))

print(result)

마무리

파이썬 개념 정리는 처음으로 글을 썼다. 백준닷컴 문제를 풀다보니 순열, 조합을 사용하는 것이 많아서 정리가 필요해서 글을 썼다. 파이썬 라이브러리를 직접 찾아보면서 글을 적다보니 시간이 많이 걸렸지만 대충 알고 있던 정보도 올바르게 알게 되어서 시간이 날 때면 라이브러리를 찾아 봐야할 것 같다.
오늘 공부한 내용 중에서 중요했던 것을 되짚어보면서 마무리 하겠다.

• itertools를 import해서 만들어져있는 순열, 조합 함수를 쓰면 쉽게 코드를 작성할 수 있다.
• 순열, 조합의 반환형은 class이기 때문에 list로 변환하면 편하게 쓸 수 있다. 하지만 필수는 아니다.
• product는 단순히 중복 순열만을 위한 함수가 아니지만 알고리즘 상 중복 순열을 만들어내는 능력이 있는 것이다.
• 함수를 제대로 공부하고 싶으면 라이브러리를 참고하자.

문제 번호 : 16938

문제 출처 : https://www.acmicpc.net/problem/16938 

Code

import itertools


def diff(num):
    return max(num) - min(num)


n, l, r, x = map(int, input().split())
a = list(map(int, input().split()))
count = 0

for i in range(2, n+1):
    numbers = list(itertools.combinations(a, i))
    for number in numbers:
        if l <= sum(number) <= r:
            if diff(number) >= x:
                count += 1

print(count)

Idea

1. N개의 난이도를 2개 이상으로 조합해서 조건에 맞는지 확인한다.
2. combination한 후 list의 요소는 iterable한 튜플 형으로 모든 난이도의 합은 sum 함수로 구현이 가능하다.
3. 차를 구하는 함수는 최대값에서 최소값을 빼주는 형태로 가능했다.
4. 두 조건을 모두 만족하면 count 값을 1씩 증가시키고 마지막에 출력한다.

이전의 문제들보다 난이도가 낮은 편이었다. 처음에는 n이 1인 경우에 print(0)을 먼저 하는 경우를 생각했었는데 print를 2번하게 되어서 채점중(99%)에서 "틀렸습니다"가 나왔다. 아무래도 마지막 입력에서 n이 1인 듯하다. n이 1이면 2개 이상을 조합으로 뽑아올 수 없다.

문제 번호 : 16922

문제 출처 : https://www.acmicpc.net/problem/16922

Code

import itertools

n = int(input())
rome = [1, 5, 10, 50]


comb = list(itertools.combinations_with_replacement(rome, n))
data = set([])
for t in comb:
    data.add(sum(t))

print(len(data))

조합(Combination) 문제라 생각하고 itertools를 사용해서 문제를 풀어봤다. rome이라는 데이터셋에서 n개의 데이터를 반복이 되도록 조합한다. comb 리스트 안에 각 원소들은 tuple 자료형이다. 그냥 조합만을 사용해서 했을 경우에는 합이 중복되는 경우가 존재했고 data라는 집합 자료형에 각 튜플들의 합을 넣어주면서 중복을 피했다.

중복이 일어나는 예시로는 n이 6인 경우 XXXXXV 와 LIIIII 이 둘 다 55의 값을 가지게 된다.

※주의할 점

• combinations_with_replacement()의 반환 형태가 class로 존재하며 이것을 iterable한 자료형으로 바꿔주어야 한다. 
• 집합 자료형은 중복을 허용하지 않고 인덱스를 가지지 않는다.

다음 코드는 itertools를 import하지 않고 반복문을 사용한 코드이다.

n = int(input())
rome = [1, 5, 10, 50]
result = set([0])
for _ in range(n):
    result = set([i + j for i in result for j in rome])

print(len(result))

처음에 n 만큼 반복하면서 result에 원소들을 넣는 것이 아니라 완전히 새로운 집합으로 만든다.

위의 for 문와 비슷하게 다음 코드를 조금 더 보기 쉽게 만든 것이다.

tmp = []
for _ in range(n):
    for i in result:
        for j in rome:
            tmp.append(i + j)
    result = set(tmp)
    tmp = []