백준 19699 소-난다! Kotlin (조합)

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

19699번: 소-난다!

문제

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

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

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

입력

첫째 줄에 농장에 있는 소들의 수 N, 선별할 소의 수 M이 주어진다.

둘째 줄에 소들의 몸무게 Hi가 주어진다.

출력

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

제한

•  1≤M≤N≤9 
•  1≤Hi≤1,000 

풀이

지금보니 순열이 아니라 조합으로 풀면 된다.

소가 2마리 있고 몸무게가 4,10이라고 할 때

순열은 (4,10), (10,4)를 모두 구하지만 조합은 (4,10)하나만 구한다.

순서만 다른 쌍을 다른 결과로 보느냐 -> 순열

순서만 다른 쌍을 같은 결과로 보느냐 -> 조합

우린 몸무게 합만 필요하므로 4+10이나 10+4나 똑같기 때문에 조합으로 풀면 된다..

아래는 순열로 푼 풀이다. 조합이 당연히 탐색 가지가 적기 때문에 더 빠르고 조합 코드는 아래 permutation 함수에서 visited와 for문이 필요 없다. 그냥 재귀에서 무조건 현재 인덱스 선택하고 다음으로 넘어가면 된다. 

 

전체적인 풀이는 다음과 같다.

1. '에라토스테네스의 체'로 소수 집합을 구해준다.

2. 순열(재귀)을 이용해 M마리의 소를 뽑는다.

3. M마리의 소를 뽑았다면 M마리 소의 몸무게 합이 소수라면 정답 리스트에 추가한다.

4. 이때 소수는 위에서 구한 소수 집합으로 O(1)만에 확인 가능하고, 정답 리스트는 set으로 두어 중복을 제거한다.

5. 정답 배열이 비어있다면 -1을 출력하고 비어있지 않다면 정렬하여 출력한다.

코드

val br = System.`in`.bufferedReader()
fun getIntList() = br.readLine().trim().split(' ').map { it.toInt() }
fun getInt() = br.readLine().trim().toInt()

val isPrime = BooleanArray(9001) { true }
val answer = mutableSetOf<Int>()
fun main() = with(System.out.bufferedWriter()) {
    //input
    val (n, m) = getIntList()
    val cows = getIntList()
    //solve
    makePrime()
    permutation(0, 0, BooleanArray(n), n, m, cows)
    //output
    write("${if (answer.isEmpty()) -1 else answer.sorted().joinToString(" ")} ")
    close()
}

fun permutation(cnt: Int, sum: Int, visited: BooleanArray, n: Int, m: Int, cows: List<Int>) {
    if (cnt == m) {
        if (isPrime[sum]) {
            //output
            answer.add(sum)
        }
        return
    }
    for (i in cows.indices) {
        if (visited[i]) continue
        visited[i] = true
        permutation(cnt + 1, sum + cows[i], visited, n, m, cows)
        visited[i] = false
    }
}

fun makePrime() {
    isPrime[1] = false
    for (i in 2..100) {
        for (j in i * 2..9000 step i) {
            isPrime[j] = false
        }
    }
}