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

백준 19699 - 소-난다!

시간 제한 1초, 메모리 제한 1024MB

# 조건

• 지난 번 헛간 청약의 당첨우(牛)가 발표됐다.
• 청약에 당첨된 소들은 날아갈 듯이 기뻐하다가 진짜로 하늘을 날았다.
• 하지만 이후로 소들은 날 수 없었다.
• 그러던 어느 날, 꿀벌에게 쏘이면 잠깐 하늘을 날 수 있다는 사실을 깨달았다.
• 이 사실이 퍼지자 소들은 다시 자유롭게 하늘을 날기 시작했다.
• 소들이 하늘을 날며 우(牛)통사고가 빈번해지자, 농부 존은 소들이 하늘을 나는 것에 제한을 두었다.
• 소들은 항의했지만 소들의 항의는 받아들여지지 않았다.
• 농장에는 $N$마리의 소가 있다. 농부 존은 소들의 몸무게의 합이 소수(prime)가 되도록 $M$마리의 소를 선별할 계획이다.
• 농부 존의 계획에 맞게 소를 선별했을 때 나올 수 있는 몸무게의 합을 모두 출력하시오.

입력

• 첫째 줄에 농장에 있는 소들의 수 N, 선별할 소의 수 M이 주어진다.
• 둘째 줄에 소들의 몸무게 Hi가 주어진다.

출력

• M마리 소들의 몸무게 합으로 만들 수 있는 모든 소수를 오름차순으로 출력하라.
• 만약 그러한 경우가 없다면 -1을 출력한다.

# 접근 방법

• 조합과 아리스토테네스의 체를 사용해준다.
• 우선 최대로 나올 수 있는 몸무게의 합이 9000이므로 9001까지의 nums 리스트를 True로 기록해준다.
• 이후 9001의 제곱 + 1까지 수를 순회하며 nums[i] = True인 경우 while 문을 이용하여 배수들을 False로 변경해준다.
• 이후 조합을 순회하며 result에 add 해준 후 출력해준다.

import sys  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  
from itertools import combinations  

N, M = map(int, input().split())  
weight = [*map(int, input().split())]  
nums = [True] * 9001  
nums[0], nums[1] = False, False  
for i in range(2, int(9001**0.5)+1):  
    if nums[i] == True:  
        j = 2  
        while i * j <= 9000:  
            nums[i*j] = False  
            j += 1  
result = set()  
for i in combinations(weight, M):  
    val = sum(i)  
    if nums[val] == True:  
        result.add(val)  
if result:  
    print(*sorted(list(result)))  
else:  
    print(-1)