[백준/BOJ] 16202번: MST 게임

문제

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

16202번: MST 게임

해설

제목에서부터 알 수 있듯이 MST를 응용하는 문제이다.

최소 신장 트리를 구하기 위해 크루스칼 알고리즘을 사용하였다.

 

문제에서 가장 중요한 문장은 아래 문장이다.

• 각 턴이 종료된 후에는 그 턴에서 구한 MST에서 가장 가중치가 작은 간선 하나를 제거한다.

결론적으로는 한 턴이 끝날 때마다 가장 작은 간선을 제거하여 k만큼 최소 신장 트리를 구하는 문제이다.

나는 이를 list를 선언해주어 해결할 수 있었다.

우선 첫번째로 크루스칼 알고리즘을 돌릴 때, 최소 신장 트리를 이루는 간선들을 list에 저장해주었다.

우선순위큐를 이용하여 가장 작은 간선부터 선택해나가므로, list의 첫번째 인덱스에는 가중치가 가장 작은 간선이 있게 된다.

첫번째 크루스칼 알고리즘이 끝났다면, removeFirst()를 통해 list의 첫번째 값을 제거해주고, 이를 다시 우선순위 큐에 넣어준다.

이를 통해 가중치가 가장 작은 간선이 제거된 우선순위큐가 되게 된다.

코드

import java.lang.StringBuilder
import java.util.*

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

private lateinit var parent: IntArray
private val list = ArrayList<Info>()

fun main() = with(System.`in`.bufferedReader()) {
    val sb = StringBuilder()
    var (nodeCnt, edgeCnt, k) = readLine().split(" ").map { it.toInt() }

    var pq = PriorityQueue<Info>()
    for(i in 1 .. edgeCnt) {
        val (start, end) = readLine().split(" ").map { it.toInt() }
        pq.add(Info(start, end, i))
    }

    var isZero = false
    for(i in 0 until k) {
        if(isZero) {
            sb.append("0 ".repeat(k - i))
            break
        }

        var getMst = kruskal(nodeCnt, pq)
        if(getMst == 0) isZero = true
        sb.append("$getMst ")

        list.removeFirst()
        pq.addAll(list)
        list.clear()
    }
    println(sb.toString())
}

private fun kruskal(n: Int, pq: PriorityQueue<Info>): Int {
    parent = IntArray(n + 1) { i -> i }
    var cost = 0
    var cnt = 0

    while(pq.isNotEmpty()) {
        val cur = pq.poll()

        if (find(cur.start) != find(cur.end)) {
            union(cur.start, cur.end)
            cost += cur.cost
            cnt++
        }
        list.add(cur)
    }

    if(cnt == n - 1) return cost
    return 0
}

private fun union(a: Int, b: Int) {
    val aRoot = find(a)
    val bRoot = find(b)

    parent[bRoot] = aRoot
}

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