[백준 1017번] [Kotlin] 소수 쌍

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

1017번: 소수 쌍

난이도 : 플래티넘 3
태그 : 수학, 정수론, 소수 판정, 에라토스테네스의 체, 이분 매칭

 

 

설명

소수 판정 + 이분 매칭 문제입니다.

이분 매칭에 관해서는 아래 글을 참고해주세요.

https://uknowblog.tistory.com/444

[알고리즘] 이분 매칭 (Bipartite Matching)

 

 

 

에라토스테네스의 체로 소수 구하기 + 소수 쌍 그래프 만들기

리스트가 있을 때, 각 수들을 짝지어 쌍의 합이 소수가 되게 하는 문제입니다.

매 번 소수인지 확인하려면 시간초과가 날게 뻔하니, 저는 에라토스테네스의 체를 이용하여

소수를 미리 구해놨습니다.

리스트의 각 원소가 1000보다 작거나 같고 중복이 없기 때문에, 최대 1999까지 해당 수가 소수인지를 체크하면 됩니다.

 

fun makePrimesArray() {
    for (i in 2..<2000) {
        if (isPrimes[i]) {
            for (j in i + i..<2000 step i) {
                isPrimes[j] = false
            }
        }
    }
}

 

var graph: Array<MutableList<Int>>

===

for (i in 0..<n) {
    for (j in i + 1..<n) {
        if (isPrimes[arr[i] + arr[j]]) {
            graph[i].add(j)
            graph[j].add(i)
        }
    }
}

 

 

graph는 리스트의 각 원소가 짝을 이뤄 소수로 나타낼 수 있는 다른 수를 담고 있습니다.

즉, graph[i][j]는 arr[i] + arr[j]가 소수임을 의미합니다.

 

모든 원소쌍을 돌며, arr[i] + arr[j]라면,

graph[i]에 j를 추가해줍니다.

반대 방향 역시 마찬가지로 간선을 추가해줍니다.

 

 

 

이분 매칭 (DFS)

val connected = ArrayList<Int>()

for (i in 0..<graph[0].size) {
    matched.fill(-1)
    var ans = 1

    matched[graph[0][i]] = 0

    for (j in 1..<n) {
        visited.fill(false)
        if (dfs(j)) ans++
    }

    if (ans == n) {
        connected.add(arr[graph[0][i]])
    }
}

==========================================

fun dfs(cur: Int): Boolean {
    if (visited[cur] || cur == 0) return false
    visited[cur] = true

    for (next in graph[cur]) {
        if (matched[next] == -1 || dfs(matched[next])) {
            matched[next] = cur
            return true
        }
    }
    return false
}

 

 

저는 리스트의 첫 번째 원소가

쌍으로 가질 수 있는(더해서 소수가 되는) 원소를 미리 fix 해놓고,

나머지 원소들에 대해 이분 매칭을 시도했습니다.

 

매칭된 개수가 n과 같다면 첫 번째 원소와 매칭된 쌍 원소를 connected에 저장합니다.

 

 

이분 매칭을 수행하는 DFS 부분은 다른 문제와 크게 다르지 않습니다.

다만 cur == 0, 즉 첫 번째 원소는 절대 양보하지 않는다는 조건을 추가해줬습니다.

 

해당 조건을 통해 첫 번째 원소를 제외하고 다른 원소들에 대해서만

짝이 최대로 되게 하는 이분 매칭 알고리즘이 수행됩니다.

 

 

 

 

소스코드

lateinit var isPrimes: BooleanArray
lateinit var graph: Array<MutableList<Int>>
lateinit var visited: BooleanArray
lateinit var matched: IntArray

fun main() = with(System.`in`.bufferedReader()) {
    val n = readLine().toInt()
    val arr = readLine().split(" ").map { it.toInt() }
    isPrimes = BooleanArray(2000) { true }
    visited = BooleanArray(n)
    matched = IntArray(n)
    graph = Array(n) { mutableListOf<Int>() }

    makePrimesArray()

    for (i in 0..<n) {
        for (j in i + 1..<n) {
            if (isPrimes[arr[i] + arr[j]]) {
                graph[i].add(j)
                graph[j].add(i)
            }
        }
    }

    val connected = ArrayList<Int>()

    for (i in 0..<graph[0].size) {
        matched.fill(-1)
        var ans = 1

        matched[graph[0][i]] = 0

        for (j in 1..<n) {
            visited.fill(false)
            if (dfs(j)) ans++
        }

        if (ans == n) {
            connected.add(arr[graph[0][i]])
        }
    }

    if (connected.isEmpty()) {
        println(-1)
        return
    }

    println(connected.sorted().joinToString(" "))
}

fun makePrimesArray() {
    for (i in 2..<2000) {
        if (isPrimes[i]) {
            for (j in i + i..<2000 step i) {
                isPrimes[j] = false
            }
        }
    }
}

fun dfs(cur: Int): Boolean {
    if (visited[cur] || cur == 0) return false
    visited[cur] = true

    for (next in graph[cur]) {
        if (matched[next] == -1 || dfs(matched[next])) {
            matched[next] = cur
            return true
        }
    }
    return false
}

 

모든 수가 다 짝지어 있을 때, 첫 번째 원소와 짝을 이룰 수 있는

connected 리스트를 정렬해 출력해줍니다.

다만, connected가 비어있다면 모든 수가 짝을 이룰 수 없는 것이므로 -1을 출력합니다.

 

 

 

후기

소수를 매 번 구하면 당연히 시간초과 나겠지...?

에라토스테네스의 체를 써야겠다! 까지는 빠르게 떠올렸으나

 

이후 어떻게 접근해야 할 지 다소 막막했던 문제였습니다.

혹시 첫 번째 원소를 미리 fix해놓고, 다른 원소에 대해서만 이분 매칭을 하면...?

그리고 이걸 첫 번째 원소와 짝을 이룰 수 있는 모든 원소에 대해 반복한다면...?

이라는 아이디어를 떠올렸고,

시간초과가 날 지도 모르나... 한 번 시도는 해보자는 생각에 제출했는데 의외로 통과했네요.