[파이썬/python] 백준 19699 - 소-난다!

문제

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

---

문제 설명

1.   농장에 N마리의 소가 있다.
2. 소들의 몸무게의 합이 소수(Prime)이 되도록 M마리의 소를 선별해야 한다.
3. 조건에 맞게 소를 선별했을 때 나올 수 있는 몸무게의 합을 모두 출력하시오

 

---

풀이

에라토스테네스의 채를 이용하여 모든 소의 몸무게의 합까지 존재하는 소수를 모두 구한다.

그 후 조합을 이용하여 나올 수 있는 몸무게의 경우의 수를 모두 구하여, 그 중 소수만을 찾는다.

1. 에라토스테네스의 채를 통해 최댓값까지 모든 소수를 찾는다.
2. combination 함수를 통해 모든 조합을 검사한다.
   1. prime[weight] == true
      해당 몸무게의 합이 소수이므로 res 배열에 넣는다.
3. res 배열을 오름차순으로 정렬한 후 출력한다.

---

코드

import sys
from itertools import combinations
input = sys.stdin.readline

MAX = 9001

def make_prime():
    prime = [i for i in range(MAX)]
    prime[1] = 0
    for i in range(2, MAX):
        for j in range(i*2, MAX, i):
            prime[j] = 0
    return prime

N,M = map(int,input().split())
cows = list(map(int,input().split()))
prime = make_prime()
res = set()
for selected in combinations(cows,M):
    weight = sum(selected)
    if prime[weight]:
        res.add(weight)
if not res:
    print(-1)
else:
    print(*sorted(res))

---

시간복잡도

• 에라토스테네스의 시간 복잡도는 $O(N*log{log_N})$이다.
  • 최악일 경우 sum(weight) = 9000이므로 $O(9*10^3 * 3) = O(10^4)$
• 조합의 시간복잡도는 $O(2^N)$
  • 최악일 경우 $N = 9$이므로 약 $O(2^9)$