[BOJ 골드4] 최소 스패닝 트리

Node data class 사용방법

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


var V = 0
var E = 0
lateinit var adj : Array<ArrayList<Node>>

data class Node(var to : Int, var weight : Int) : Comparable<Node>{
    override fun compareTo(other: Node): Int {
       return this.weight - other.weight
    }

}

fun main() = with(System.`in`.bufferedReader()) {
    val (n,m) = readLine().split(" ").map{it.toInt()}
    V = n
    E = m
    adj  = Array(n+1){
        ArrayList()
    }
    for(i in 1..m){
        val (a,b,c) = readLine().split(" ").map{it.toInt()}
        adj[a].add(Node(b,c))
        adj[b].add(Node(a,c))
    }
    println(prim())
}

fun prim() : Long{
    val visited = Array(V+1){false}
    val pq = PriorityQueue<Node>()

    pq.add(Node(1,0))
    var ans = 0L
    var cnt = 0
    while(!pq.isEmpty()){
        val temp = pq.poll()
        if(!visited[temp.to]){
            ans += temp.weight
            visited[temp.to] = true
            if(++cnt == V) return ans
            for(node in adj[temp.to]){
                if(!visited[node.to]){
                    pq.add(node)
                }
            }
        }
    }
    return ans
}

Pair 사용 방법

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


var V = 0
var E = 0
lateinit var adj : Array<ArrayList<Pair<Int,Int>>>


fun main() = with(System.`in`.bufferedReader()) {
    val (n,m) = readLine().split(" ").map{it.toInt()}
    V = n
    E = m
    adj  = Array(n+1){
        ArrayList()
    }
    for(i in 1..m){
        val (a,b,c) = readLine().split(" ").map{it.toInt()}
        adj[a].add(Pair(b,c))
        adj[b].add(Pair(a,c))
    }
    println(prim())
}


fun prim() : Long {
    var ans = 0L
    var cnt = 0
    var pq = PriorityQueue<Pair<Int,Int>>{a,b -> a.second-b.second}
    pq.add(Pair(1,0))
    var visit = Array(V+1){false}
    while(pq.isNotEmpty()){
        var node= pq.poll()
        if(!visit[node.first]){
            ans += node.second
            if(++cnt==V) return ans
            visit[node.first] = true
            for(next in adj[node.first]){
                if(!visit[next.first]){
                    pq.add(next)
                }
            }
        }
    }
    return ans
}

시간 절약을 위해선 Pair를 쓰는게 더 맞을듯 하다

좋은 웹페이지 즐겨찾기