[백준] 1774번: 우주신과의 교감 - Kotlin[코틀린]

문제

 

1774번: 우주신과의 교감

 

---

풀이

 

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

 

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

 

---

코드

 

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

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

// 간선을 나타내는 데이터 클래스
data class Edge(val start: Int, val end: Int, val value: Double) : Comparable<Edge> {
    override fun compareTo(other: Edge): Int = value.compareTo(other.value)
}

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

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

    // 입력
    val st = StringTokenizer(br.readLine())
    val num = st.nextToken().toInt() // 별의 개수
    val root = st.nextToken().toInt() // 초기에 선택된 별의 개수

    // 별의 좌표 정보 저장
    val list = ArrayList<Star>(num)
    // 간선 정보 저장을 위한 리스트
    val graph = ArrayList<Edge>(num)

    repeat(num) {
        with(StringTokenizer(br.readLine())) {
            list.add(Star(nextToken().toInt(), nextToken().toInt(), it + 1))
        }
    }

    // 모든 별 간의 거리를 계산하여 간선 정보 저장
    for (i in 0 until num) {
        for (j in i + 1 until num) {
            graph.add(Edge(list[i].index, list[j].index, distance(list[i], list[j])))
        }
    }
    graph.sort() // 거리 순으로 정렬

    // 유니온-파인드를 위한 부모 배열 초기화
    val parent = IntArray(num + 1) { it }

    // 유니온-파인드에서 루트 찾기 함수
    fun find(a: Int) : Int {
        return if(parent[a]==a) a
        else find(parent[a])
    }
    // 두 집합 합치기 함수
    fun union(a: Int, b: Int) {
        val x = find(a)
        val y = find(b)
        if (x != y) {
            parent[y] = x
        }
    }

    // 크루스칼 알고리즘을 이용한 최소 스패닝 트리 계산 함수
    fun kruskal(): Double {
        var ans = 0.0

        // 간선을 하나씩 확인하며 사이클이 생기지 않으면 해당 간선을 선택
        for (edge in graph) {
            if (find(edge.start) != find(edge.end)) {
                ans += edge.value
                union(edge.start, edge.end)
            }
        }

        return ans
    }

    // 초기 선택된 별들을 하나의 집합으로 합치기
    repeat(root) {
        with(StringTokenizer(br.readLine())) {
            union(nextToken().toInt(), nextToken().toInt())
        }
    }

    // 결과 출력
    bw.write(String.format("%.2f", kruskal()))
    bw.flush()
    bw.close()
    br.close()
}