109 lines
2.8 KiB
Swift
109 lines
2.8 KiB
Swift
struct MinHeap<T> {
|
|
var nodes: [T] = []
|
|
let sort: (T, T) -> Bool
|
|
|
|
init(sort: @escaping (T, T) -> Bool) {
|
|
self.sort = sort
|
|
}
|
|
|
|
var isEmpty: Bool {
|
|
return nodes.isEmpty
|
|
}
|
|
|
|
mutating func insert(_ element: T) {
|
|
nodes.append(element)
|
|
shiftUp(from: nodes.count - 1)
|
|
}
|
|
|
|
mutating func pop() -> T? {
|
|
if nodes.isEmpty { return nil }
|
|
if nodes.count == 1 { return nodes.removeLast() }
|
|
|
|
let root = nodes[0]
|
|
nodes[0] = nodes.removeLast()
|
|
shiftDown(from: 0)
|
|
return root
|
|
}
|
|
|
|
private mutating func shiftUp(from index: Int) {
|
|
var child = index
|
|
var parent = (child - 1) / 2
|
|
|
|
while child > 0 && sort(nodes[child], nodes[parent]) {
|
|
nodes.swapAt(child, parent)
|
|
child = parent
|
|
parent = (child - 1) / 2
|
|
}
|
|
}
|
|
|
|
private mutating func shiftDown(from index: Int) {
|
|
var parent = index
|
|
|
|
while true {
|
|
let left = parent * 2 + 1
|
|
let right = parent * 2 + 2
|
|
var candidate = parent
|
|
|
|
if left < nodes.count && sort(nodes[left], nodes[candidate]) {
|
|
candidate = left
|
|
}
|
|
if right < nodes.count && sort(nodes[right], nodes[candidate]) {
|
|
candidate = right
|
|
}
|
|
|
|
if candidate == parent { return }
|
|
|
|
nodes.swapAt(parent, candidate)
|
|
parent = candidate
|
|
}
|
|
}
|
|
}
|
|
|
|
if let input1 = readLine(),
|
|
let nums1 = input1.split(separator: " ").compactMap({Int($0)}) as? [Int],
|
|
let V = nums1.first, let E = nums1.last
|
|
{
|
|
var graph: [[(node: Int, w: Int)]] = Array(repeating: [], count: V+1)
|
|
for _ in 1...E {
|
|
guard let input2 = readLine(),
|
|
let nums2 = input2.split(separator: " ").compactMap({Int($0)}) as? [Int],
|
|
nums2.count == 3
|
|
else { break }
|
|
|
|
let (u, v, w) = (nums2[0], nums2[1], nums2[2])
|
|
graph[u].append((node: v, w: w))
|
|
graph[v].append((node: u, w: w))
|
|
}
|
|
|
|
var pq = MinHeap<(node: Int, w: Int)>{$0.w < $1.w}
|
|
var visited = Array(repeating: false, count: V+1)
|
|
|
|
pq.insert((node: 1, w: 0))
|
|
var count = 0
|
|
var totalWeight = 0
|
|
|
|
while !pq.isEmpty {
|
|
guard let current = pq.pop() else {break}
|
|
|
|
if visited[current.node] {
|
|
continue
|
|
}
|
|
|
|
visited[current.node] = true
|
|
count += 1
|
|
totalWeight += current.w
|
|
|
|
if count == V {
|
|
break
|
|
}
|
|
|
|
for next in graph[current.node] {
|
|
if !visited[next.node] {
|
|
pq.insert(next)
|
|
}
|
|
}
|
|
}
|
|
|
|
print(totalWeight)
|
|
}
|