[백준/BOJ] 6497번: 전력난

문제

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

6497번: 전력난

해설

문제를 보면 MST 알고리즘을 사용해야 한다는 것을 알아야 한다.

최소 신장 트리를 구해 총 비용에서 최소 비용을 빼어 절약할 수 있는 최대 액수를 구한다.

어떠한 정점에서 시작해야 된다와 같은 조건이 없으므로, 크루스칼 알고리즘을 사용하였다.

 

이 문제에서 주의할 점은 입력에 있다.

처음 제출했을 때, 런타임 에러가 나 무엇이 문제인지 고민하다 질문 검색에 들어가 나랑 같은 케이스인 사람들이 있는 것을 보았다. 런타임 에러가 났던 이유는 "입력은 여러 개의 테스트 케이스로 구분되어 있다."에 있었다.

입력의 마지막에 0 0이 오는 것이 아닌 여러 개의 테스트 케이스를 받다가, 집의 수와 길의 수가 0일 때 종료하는 문제였다.

소스 코드

import java.util.*

private lateinit var parent: IntArray

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

fun main() = with(System.`in`.bufferedReader()) {
    while(true) {
        val (homeCnt, roadCnt) = readLine().split(" ").map { it.toInt() }
        if(homeCnt == 0 && roadCnt == 0) break
        var totalCost = 0

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

        val pq = PriorityQueue<Info>()
        repeat(roadCnt) {
            val (start, end, cost) = readLine().split(" ").map { it.toInt() }
            totalCost += cost
            pq.add(Info(start, end, cost))
        }

        var minCost = 0
        var cnt = 0
        while (cnt != homeCnt - 1) {
            val cur = pq.poll()
            if (find(cur.start) != find(cur.end)) {
                union(cur.start, cur.end)
                minCost += cur.cost
                cnt++
            }
        }

        println(totalCost - minCost)
    }
}

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

    parent[bParent] = aParent
}

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