[백준] 4386번: 별자리 만들기 - Kotlin[코틀린]

문제

 

4386번: 별자리 만들기

 

---

풀이

 

 최소 스패닝 트리를 구하는 문제로 프림과 크루스칼 알고리즘을 이용해 풀이할 수 있다. 이번 문제에서는 프림 알고리즘을 사용하였다.

 

 문제는 각 정점의 위치가 주어지기 때문에, 정점들 간의 거리를 계산하여 간선의 정보를 우선 순위 큐 리스트에 저장해준다. 그리고 나서 프림 알고리즘을 이용하여 최소 스패닝 트리를 구하고 결과를 출력해준다.

 

---

코드

 

import java.util.*
import kotlin.math.*

// 별의 좌표 정보를 저장하는 데이터 클래스
data class Star(val x: Double, val y: Double, val index: Int)

// 우선순위 큐에서 사용할 간선 정보를 저장하는 데이터 클래스
data class Node(val index: Int, val value: Double): Comparable<Node> {
    override fun compareTo(other: Node): Int = value.compareTo(other.value)
}

// 두 별 간의 거리를 계산하는 함수
fun distance(start: Star, end: Star) = sqrt((start.x - end.x).pow(2) + (start.y - end.y).pow(2))

fun main() {
    val br = System.`in`.bufferedReader()
    val bw = System.out.bufferedWriter()

    // 별의 개수 입력
    val num = br.readLine().toInt()

    // 별의 좌표와 간선 정보를 저장할 리스트 초기화
    val list = ArrayList<Star>(num)
    val graph = ArrayList<PriorityQueue<Node>>(num)
    val visited = BooleanArray(num)

    // 별의 좌표 입력 및 간선 리스트 초기화
    repeat(num) {
        graph.add(PriorityQueue())
        with(StringTokenizer(br.readLine())) {
            list.add(Star(nextToken().toDouble(), nextToken().toDouble(), it))
        }
    }

    // 모든 별 간의 거리를 계산하여 우선순위 큐에 추가
    for (start in list) {
        for (end in list) {
            if (end == start) continue
            graph[start.index].add(Node(end.index, distance(start, end)))
        }
    }

    // 프림 알고리즘을 이용하여 최소 스패닝 트리의 가중치 합 계산
    fun prim(): Double {
        val queue = PriorityQueue<Node>()
        var sum = 0.0
        var cnt = 0

        // 임의의 시작 노드를 큐에 추가
        queue.add(Node(0, 0.0))

        while (queue.isNotEmpty()) {
            val tmp = queue.poll()

            // 이미 방문한 노드라면 스킵
            if (visited[tmp.index]) continue

            // 현재 노드 방문 처리 및 가중치 합산
            visited[tmp.index] = true
            sum += tmp.value

            // 트리에 연결된 미방문 노드들을 큐에 추가
            for (next in graph[tmp.index]) {
                if (!visited[next.index]) {
                    queue.add(next)
                }
            }

            // 트리의 노드 개수 증가
            if (++cnt == num) return sum
        }

        return sum
    }

    // 결과 출력
    bw.write("${prim()}")
    bw.flush()
    bw.close()
    br.close()
}