[백준/BOJ] 2887번: 행성 터널

문제

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

2887번: 행성 터널

해설

모든 행성을 연결했을 때, 필요한 최소 비용을 구하는 문제이므로 크루스칼 알고리즘을 사용하였다.

처음 도전했을 때는 행성 간의 거리를 이중 반복문을 이용해 모두 계산하여 풀었다.

하지만 이렇게 할 시 이중 반복문에서 O(n²)이고, n의 범위가 10만이기 때문에 메모리 초과가 발생하였다.

 

이를 해결하기 위해 모든 행성을 비교하지 않고 두 행성간 가장 작은 거리를 우선순위 큐에 넣어주었다.

즉, x, y, z좌표를 오름차순으로 정렬한 후, 두 행성의 x축 혹은 y축 혹은 z축의 차와 좌표를 우선순위 큐에 넣어주었다.

우선순위 큐의 기준을 거리로 잡아 가장 짧은 비용이 드는 두 행성을 연결해나가면, 결국 모든 행성이 연결되게 된다.

소스 코드

import java.util.*
import kotlin.collections.ArrayList

data class Info(val from: Int, val to: Int, val dist: Int): Comparable<Info> {
    override fun compareTo(other: Info): Int = dist - other.dist
}

private lateinit var parent: IntArray

fun main() = with(System.`in`.bufferedReader()) {
    val n = readLine().toInt()

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

    val posX = ArrayList<Pair<Int, Int>>()
    val posY = ArrayList<Pair<Int, Int>>()
    val posZ = ArrayList<Pair<Int, Int>>()
    repeat(n) { i ->
        val (x, y, z) = readLine().split(" ").map { it.toInt() }
        posX.add(Pair(i, x))
        posY.add(Pair(i, y))
        posZ.add(Pair(i, z))
    }

    val sortedX = posX.sortedBy { it.second }
    val sortedY = posY.sortedBy { it.second }
    val sortedZ = posZ.sortedBy { it.second }

    val pq = PriorityQueue<Info>()
    for(i in 0 until n - 1) {
        pq.add(Info(sortedX[i].first, sortedX[i+1].first, sortedX[i+1].second - sortedX[i].second))
        pq.add(Info(sortedY[i].first, sortedY[i+1].first, sortedY[i+1].second - sortedY[i].second))
        pq.add(Info(sortedZ[i].first, sortedZ[i+1].first, sortedZ[i+1].second - sortedZ[i].second))
    }

    var ans = 0L
    while(pq.isNotEmpty()) {
        val cur = pq.poll()

        if(find(cur.from) != find(cur.to)) {
            union(cur.from, cur.to)
            ans += cur.dist
        }
    }

    println(ans)
}

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

    parent[bParent] = aParent
}

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