[백준/BOJ] 14621번: 나만 안되는 연애

문제

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

14621번: 나만 안되는 연애

해설

연결그래프에서 최소 신장 트리를 찾는 문제이다.

고려할 점은 다음과 같다.

1. 연결된 두 노드가 남초, 여초 학교여야 한다.
   두 노드가 모두 남초거나, 여초라면 연결되지 않는다.
   만약 두 노드가 같다면, 우선순위 큐에 넣을 필요가 없다.
2. 두 노드간 한 개의 간선만 있는 것이 아니라 여러 개의 간선이 있을 수 있고, 이들 간 거리는 다를 수 있다.
   이는 우선순위큐를 사용할 것이기에 신경쓰지 않아도 된다.
   짧은 비용부터 계산을 해주면 되기 때문이다
3. 모든 노드(학교)가 연결되지 않는다면 -1을 출력한다.
   연결된 노드의 개수를 저장할 변수를 선언해준다.
   마지막에 노드의 개수가 (학교의 개수 - 1)개라면 모든 노드가 연결된 것이다.

위의 3가지를 제외하면 크루스칼 알고리즘의 기본을 따라가면 된다.

크루스칼 알고리즘에 대해 잘 모른다면 크루스칼 알고리즘을 참고하면 좋을 것이다.

소스 코드

import java.util.*

private lateinit var parent: IntArray

data class Connect(val start: Int, val end: Int, val cost: Int): Comparable<Connect> {
    override fun compareTo(other: Connect): Int = cost - other.cost
}

fun main() = with(System.`in`.bufferedReader()) {
    val (nodeCnt, edgeCnt) = readLine().split(" ").map { it.toInt() }
    parent = IntArray(nodeCnt + 1)
    repeat(nodeCnt + 1) { i -> parent[i] = i }

    val boyOrGirl = readLine().split(" ")
    val section = Array(nodeCnt + 1 ) { "" }
    for(i in 1..nodeCnt) section[i] = boyOrGirl[i-1]

    val pq = PriorityQueue<Connect>()
    repeat(edgeCnt) {
        val (start, end, cost) = readLine().split(" ").map { it.toInt() }
        if(section[start] != section[end]) pq.add(Connect(start, end, cost))
    }

    var ans = 0
    var cnt = 0
    while(pq.isNotEmpty()) {
        val cur = pq.poll()
        if(find(cur.start) != find(cur.end)) {
           union(cur.start, cur.end)
           ans += cur.cost
            cnt++
        }
    }

    if(cnt != nodeCnt - 1) println(-1)
    else println(ans)
}

private fun union(a: Int, b: Int) {
    val aParent = find(a)
    val bParent = find(b)

    parent[bParent] = aParent
}

private fun find(i: Int): Int {
    if(parent[i] != i) parent[i] = find(parent[i])
    return parent[i]
}