→ Problems

[Algorithm] 백준 - 1197번 최소 스패닝 트리 (Swift)

Swift librarian 2024. 6. 18. 10:16

문제 소개

 

문제 풀이

최소 신장 트리 (Minimum Spanning Tree)

그래프 상에서 모든 노드가 사이클 없이 연결되어 있고, 정점이 최소 간선의 합으로 연결된 부분 그래프이다.

출처: 위키피디아

이런 최소 신장 트리는 도로 건설 문제, 네트워크 라우팅 경로, 통신에서 전화선 길이에 사용되는 등 실제 자주 활용되는 자료구조라고 한다. 위 문제는 주어진 간선의 정보를 가지고 최소 신장 트리를 구하는 문제였다.
 

크루스칼 알고리즘

최소 신장 트리를 찾는 알고리즘에 크루스칼 알고리즘이 있다. 주어진 간선을 오름차순 정렬 후 간선을 하나씩 확인하며 간선간 사이클이 발생하지 않을 경우 최소신장트리에 포함시키는 방법이다. 사이클 발생 여부는 union find 알고리즘을 활용한다.

import Foundation

let ve = readLine()!.split(separator: " ").map { Int($0)! }
let (v, e) = (ve[0], ve[1])
var parents = Array(0...v)
var edges = [[Int]]()

func getParent(_ n: Int) -> Int {
    if parents[n] == n { return n }
    return getParent(parents[n])
}

func unionParent(_ a: Int, _ b: Int) {
    let a = getParent(a)
    let b = getParent(b)
    
    if a < b {
        parents[b] = a
    } else {
        parents[a] = b
    }
}

for _ in 0..<e {
    let abc = readLine()!.split(separator: " ").map { Int($0)! }
    let (a, b, c) = (abc[0], abc[1], abc[2])
    edges.append([a, b, c])
}

edges.sort { $0[2] < $1[2] }

var answer = 0

for edge in edges {
    let a = edge[0]
    let b = edge[1]
    let cost = edge[2]
    
    if getParent(a) != getParent(b) {
        unionParent(a, b)
        answer += cost
    }
}

print(answer)