BOJ 19699 - 소-난다! (Python)

S2, 백트랙킹

문제

지난 번 헛간 청약의 당첨우(牛)가 발표됐다. 청약에 당첨된 소들은 날아갈 듯이 기뻐하다가 진짜로 하늘을 날았다. 하지만 이후로 소들은 날 수 없었다. 그러던 어느 날, 꿀벌에게 쏘이면 잠깐 하늘을 날 수 있다는 사실을 깨달았다. 이 사실이 퍼지자 소들은 다시 자유롭게 하늘을 날기 시작했다.

소들이 하늘을 날며 우(牛)통사고가 빈번해지자, 농부 존은 소들이 하늘을 나는 것에 제한을 두었다. 소들은 항의했지만 소들의 항의는 받아들여지지 않았다.

농장에는 $N$마리의 소가 있다. 농부 존은 소들의 몸무게의 합이 소수(prime)가 되도록 $M$마리의 소를 선별할 계획이다. 농부 존의 계획에 맞게 소를 선별했을 때 나올 수 있는 몸무게의 합을 모두 출력하시오.

풀이

• 백트랙킹 실버 단계까지는 문제 푸는 유형이 거의 비슷하다.
• N과M 시리즈만 많이 풀어봤어도 어려움없이 해결 가능

from sys import stdin
import math

n, m = map(int, stdin.readline().split())
arr = sorted(list(map(int ,stdin.readline().split())))

visited = [0] * n
res = []
total = set()

def isPrime(num):
    for i in range(2, int(math.sqrt(num) + 1)):
        if num % i == 0:
            return False
    return True

def dfs(depth):
    if depth == m:
        if isPrime(sum(res)) and sum(res) not in total:
            total.add(sum(res))
        return

    for i in range(n):
        if not visited[i]:
            res.append(arr[i])
            visited[i] = 1
            dfs(depth + 1)
            visited[i] = 0
            res.pop()

dfs(0)

if len(total) == 0:
    print("-1")
else:
    print(*sorted(list(total)))