백준 14621 나만 안되는 연애 Kotlin (최소 스패닝 트리)

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

14621번: 나만 안되는 연애

문제

깽미는 24살 모태솔로이다. 깽미는 대마법사가 될 순 없다며 자신의 프로그래밍 능력을 이용하여 미팅 어플리케이션을 만들기로 결심했다. 미팅 앱은 대학생을 타겟으로 만들어졌으며 대학교간의 도로 데이터를 수집하여 만들었다.

이 앱은 사용자들을 위해 사심 경로를 제공한다. 이 경로는 3가지 특징을 가지고 있다.

1. 사심 경로는 사용자들의 사심을 만족시키기 위해 남초 대학교와 여초 대학교들을 연결하는 도로로만 이루어져 있다.
2. 사용자들이 다양한 사람과 미팅할 수 있도록 어떤 대학교에서든 모든 대학교로 이동이 가능한 경로이다.
3. 시간을 낭비하지 않고 미팅할 수 있도록 이 경로의 길이는 최단 거리가 되어야 한다.

만약 도로 데이터가 만약 왼쪽의 그림과 같다면, 오른쪽 그림의 보라색 선과 같이 경로를 구성하면 위의 3가지 조건을 만족하는 경로를 만들 수 있다.

이때, 주어지는 거리 데이터를 이용하여 사심 경로의 길이를 구해보자.

입력

입력의 첫째 줄에 학교의 수 N와 학교를 연결하는 도로의 개수 M이 주어진다. (2 ≤ N ≤ 1,000) (1 ≤ M ≤ 10,000)

둘째 줄에 각 학교가 남초 대학교라면 M, 여초 대학교라면 W이 주어진다.

다음 M개의 줄에 u v d가 주어지며 u학교와 v학교가 연결되어 있으며 이 거리는 d임을 나타낸다. (1 ≤ u, v ≤ N) , (1 ≤ d ≤ 1,000)

출력

깽미가 만든 앱의 경로 길이를 출력한다. (모든 학교를 연결하는 경로가 없을 경우 -1을 출력한다.)

알고리즘 분류

• 그래프 이론
• 최소 스패닝 트리

풀이

최소 스패닝 트리 문제이다.

최소 스패닝 트리는 크루스칼 or 프림 알고리즘으로 풀 수 있고 본인은 크루스칼로 풀었다.

문제의 2번 3번 조건이 최소 스패닝 트리를 구하라는 조건이다.

2번 (모든 노드를 연결해라)
3번 (최소 비용으로 연결해라)

최소 스패닝 트리에 1번 조건이 추가되었는데, 연결하려는 두 노드가 성별이 같을 때만 연결해 주면 된다.(최소 비용의 간선일 때)

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

1.주어진 간선을 가중치 오름차순으로 정렬한다.

2. 성별이 다른 경우, 두 노드가 아직 연결되지 않은 경우 unionParent로 연결해준다.

3. 전체 간선을 한 번 쫙 돌면 최소 비용의 간선만 사용하여 모든 노드가 연결된다.

4. 하지만 이 문제의 입력에선 모든 노드가 연결되지 않는 경우도 있다. 성별이라는 조건이 있기 때문.

5. 따라서 크루스칼 로직을 끝내고 모든 노드가 연결되어있는지 (같은 부모를 가지고 있는지) 확인해서 -1 혹은 answer를 출력한다.

코드

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

var n = 0
var m = 0
lateinit var parent: IntArray

fun getParent(x: Int): Int {
    return if (parent[x] == x) x else getParent(parent[x]).also { parent[x] = it }
}

fun unionParent(x: Int, y: Int) {
    val xx = getParent(x)
    val yy = getParent(y)
    if (xx < yy) {
        parent[yy] = xx
    } else {
        parent[xx] = yy
    }
}

fun isSameParent(x: Int, y: Int) = getParent(x) == getParent(y)

fun main() = with(System.out.bufferedWriter()) {
    //input
    getIntList().also {
        n = it[0]
        m = it[1]
    }
    val sex = br.readLine().split(" ")
    val input = Array(m){ getIntList()}
    var answer = 0
    parent = IntArray(n + 1) { it }
    input.sortBy { it[2] }

    //Kruskal
    for ((u, v, d) in input) {
        //성별이 다른 경우만 (1번 조건)
        if (sex[u - 1] != sex[v - 1]) {
            if (!isSameParent(u, v)) {
                unionParent(u, v)
                answer += d
            }
        }
    }
    //output
    //하나라도 연결되지 않은 경우 -1
    val first = getParent(1)
    for(i in 2 .. n){
        if(first != getParent(i)){
            write("-1")
            close()
            return
        }
    }
    write("$answer")
    close()
}