Cut property를 이용.
import Foundation
var dict = [Int: [[Int]]]()
func solution(_ n:Int, _ costs:[[Int]]) -> Int {
for cost in costs {
if let _ = dict[cost[0]] {
dict[cost[0]]!.append( [cost[1], cost[2]])
} else {
dict[cost[0]] = [[cost[1], cost[2]]]
}
if let _ = dict[cost[1]] {
dict[cost[1]]!.append([cost[0], cost[2]])
} else {
dict[cost[1]] = [[cost[0], cost[2]]]
}
}
return prim(len: n)
}
func prim(len: Int) -> Int {
var f = Array(repeating: false, count: len)
var pq = [Int]()
var cost = Array(repeating: Int.max, count: len)
var e = Array(repeating: -1, count: len)
//
var t = [[Int]]()
for v in 0..<len {
pq.append(v)
}
while !pq.isEmpty {
pq.sort { cost[$0] > cost[$1] }
let v = pq.removeLast()
f[v] = true
if e[v] != -1 {
print(e[v], v)
t.append([e[v], v])
}
for next in dict[v]! {
if f[next[0]] == false && next[1] < cost[next[0]] {
cost[next[0]] = next[1]
e[next[0]] = v
}
}
}
var ans = 0
for edge in t {
for next in dict[edge[0]]! {
if next[0] == edge[1] {
ans += next[1]
break
}
}
}
//print(t)
return ans
}