[백준/BOJ] 10423번: 전기가 부족해

문제

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

10423번: 전기가 부족해

해설

크루스칼 알고리즘을 응용하여 해결할 수 있는 문제이다.

발전소의 개수가 한개라면, 기본적인 크루스칼 알고리즘으로 최소 비용을 구할 수 있다.

하지만 이 문제에서는 발전소의 개수가 여러 개일 수도 있다.

또한 모든 간선이 연결되어야 하는 것이 아니고, 발전소를 통해 전기를 공급받기만 하면 된다.

나는 미리 발전소의 위치를 -1로 해주었다.

이처럼 해준 이유는 발전소와 연결된 노드들의 부모 노드는 -1이 되며, union 과정에서 발전소끼리의 연결을 피할 수 있다.

모든 노드의 부모 노드가 -1이 되면 종료한다.

소스 코드

import java.util.*

private lateinit var parent: IntArray

data class Cable(val firstCity: Int, val secondCity: Int, val cost: Int): Comparable<Cable> {
    override fun compareTo(other: Cable): Int = cost - other.cost
}

fun main() = with(System.`in`.bufferedReader()) {
    val (cityCnt, cableCnt, powerCnt) = readLine().split(" ").map { it.toInt() }

    parent = IntArray(cityCnt + 1)
    repeat(cityCnt + 1) { i -> parent[i] = i }

    val st = StringTokenizer(readLine())
    repeat(powerCnt)  {
        parent[st.nextToken().toInt()] = -1
    }

    val pq = PriorityQueue<Cable>()
    repeat(cableCnt) {
        val (firstCity, secondCity, cost) = readLine().split(" ").map { it.toInt() }
        pq.add(Cable(firstCity, secondCity, cost))
    }

    var ans = 0
    while(pq.isNotEmpty()) {
        val cur = pq.poll()
        if(find(cur.firstCity) != find(cur.secondCity)) {
            union(cur.firstCity, cur.secondCity)
            ans += cur.cost

            if(allConnect()) break
        }
    }

    println(ans)
}

private fun allConnect(): Boolean {
    repeat(parent.size) { i ->
        if(parent[i] != -1) return false
    }
    return true
}

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

    if(aParent != bParent) {
        if(aParent == -1) parent[bParent] = aParent
        else if(bParent == -1) parent[aParent] = bParent
        else {
            parent[bParent] = aParent
        }
    }
}

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