[Python][백준 BOJ Bronze I] 2309. 일곱 난쟁이 : 순열과 조합, combinations, permutations, product, itertools

itertools

itertools는 파이썬에서 반복되는 데이터를 처리하는 기능을 포함하고 있는 라이브러리이다.

코딩 테스트에서 유용하게 사용할 수 있는 permutations, combinations 클래스 등을 제공한다.

permutations

리스트와 같은 iterable 객체에서 r개의 데이터를 뽑아 일렬로 나열하는 모든 경우(= 순열)를 계산하는 클래스이다. 순열은 순서가 있는 조합. 서로 다른 n개에서 r개를 택하여 일렬로 줄을 세우는 경우의 수이다.

permutations는 클래스이기 때문에, 객체 초기화 이후에는 리스트 자료형으로 변환하여 사용한다.

# 리스트에서 2개(r=2)를 뽑아 나열하는 모든 경우의 수
from itertools import permutations

data = ['A', 'B', 'C']
result = list(permutations(data, 3))    # data에서 2개를 뽑아 나열하는 모든 순열
print(result)
# [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

combinations

리스트와 같은 iterable 객체에서 r개의 데이터를 뽑아 순서를 고려하지 않고 나열하는 모든 경우(= 조합)를 계산하는 클래스이다. 조합은 서로 다른 n개에서 r개를 택하는 경우의 수이다.

combinations는 클래스이기 때문에, 객체 초기화 이후에는 리스트 자료형으로 변환하여 사용한다.

# 리스트에서 2개(r=2)를 뽑아 순서에 상관없이 나열하는 모든 경우의 수 출력
from itertools import combinations

data = ['A', 'B', 'C']
result = list(combinations(data, 2))    # data에서 2개를 뽑는 모든 조합
print(result)
# [('A', 'B'), ('A', 'C'), ('B', 'C')]

product

permutations와 같이 리스트와 같은 iterable 객체에서 r개의 데이터를 뽑아 일렬로 나열하는 모든 경우(= 순열)를 계산하나, 원소를 중복하여 뽑는다. (= 중복 순열)

product 객체를 초기화할 때는 데이터의 수를 repeat 속성값으로 넣어준다. 

product 역시 클래스이기 때문에, 객체 초기화 이후에는 리스트 자료형으로 변환하여 사용한다.

# 리스트에서 중복을 포함하여 2개(r=2)를 뽑아 나열하는 모든 경우 출력
from itertools import product
data = ['A', 'B', 'C']
result = list(product(data, repeat=2))    # data에서 2개를 뽑는 모든 중복 순열
print(result)
# [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'B'), ('B', 'C'), ('C', 'A'), ('C', 'B'), ('C', 'C')]

+) iterable 객체란?

반복 가능한 객체.

대표적으로는 list, dict, set, str, bytes, tuple, range가 있다.

 

---

2309. 일곱 난쟁이

아래 문제를 풀 때 combinations 클래스를 사용할 수 있다 ! 사실 이 문제를 풀면서 combinations에 대해 알게 되었다.

 

BOJ 2309. 일곱 난쟁이

1. 아홉 난쟁이의 키를 모두 합하고 (sum)
2. 아홉 난쟁이 중 두 명을 뽑아서
3. sum에서 두 명의 키를 뺐을 때 100이 되면
4. 그 두명을 제외한 나머지 일곱 명의 키를 출력

위의 순서로 문제를 풀어야겠다고 생각했다.

combinations에 대해 몰랐을 땐, 두 난쟁이를 뽑기 위해 이중 for문을 사용해 코드를 작성했다.

Solution 1 - 이중 for문

import sys
n = [int(sys.stdin.readline().rstrip()) for i in range(9)]

sum = sum(n)
remove_one = 0
remove_two = 0
# 2개의 요소를 뻈을 떄 100이 나오면 -> n에서 두개 요소 삭제
for j in range(0, len(n) - 1):
    for k in range(j + 1, len(n)):
        if sum - n[j] - n[k] == 100:
            # n.remove(n[j])
            # n.remove(n[k]) => 위에서 n[j] 제거하면 n[k] 달라질 수 있음
            remove_one = n[j]
            remove_two = n[k]
n.remove(remove_one)
n.remove(remove_two)

n.sort()
for l in n:
    print(l)

처음 풀었던 방식.

 

제외할 두 명을 뽑기 위해 이중 for문을 사용했고, 사람이기 때문에 j와 k가 중복되지 않도록 range를 설정했다.

그러나 조합을 사용하니 더욱 간단하게 제외할 두 명을 뽑을 수 있었다!

Solution 2 - combinations

import sys
from itertools import combinations
n = [int(sys.stdin.readline().rstrip()) for _ in range(9)]

remove_one = 0
remove_two = 0
# n에서 합이 sum - 100 인 두 요소 찾기
for i in combinations(n, 2):
    if sum(n) - sum(i) == 100:
        remove_one, remove_two = i
        break
n.remove(remove_one)
n.remove(remove_two)

n.sort()
for l in n:
    print(l)

위에서 푼 방식대로 combinations을 사용해 작성한 코드이다.

이중 for문으로 구현할 땐 제외할 2개의 요소를 구하는 게 일곱 난쟁이에 포함할 7개의 요소를 구하는 것보다 더 쉽지만, 조합을 사용한다면 합이 100인 7개의 요소를 구하는 게 어려운 일이 아니기 때문에 아래의 코드처럼 다시 풀어보았다.

Solution 3 - combinations

import sys
from itertools import combinations
n = [int(sys.stdin.readline().rstrip()) for _ in range(9)]

# n에서 합이 100 인 일곱 개의 요소 찾기
for i in combinations(n, 7):
    if sum(i) == 100:
        for j in sorted(i):
            print(j)
        break

처음에 비해 많이 짧아졌다! 또한 시간 복잡도도 개선되었다.

 

solution 1의 시간 복잡도는 O(n³), solution 2의 시간 복잡도는 O(n²), solution 3의 시간 복잡도는 O(1)이다.

점점 좋은 코드를 작성했다는 점에서 매우매우매우 뿌듯 🥰

---

https://www.acmicpc.net/problem/2309

2309번: 일곱 난쟁이