[백준] 1197번 최소 스패닝 트리 [Swift]

문제 링크

1197번: 최소 스패닝 트리

문제

그래프가 주어졌을 때, 그 그래프의 최소 스패닝 트리를 구하는 프로그램을 작성하시오.

최소 스패닝 트리는, 주어진 그래프의 모든 정점들을 연결하는 부분 그래프 중에서 그 가중치의 합이 최소인 트리를 말한다.

입력

첫째 줄에 정점의 개수 V(1 ≤ V ≤ 10,000)와 간선의 개수 E(1 ≤ E ≤ 100,000)가 주어진다. 다음 E개의 줄에는 각 간선에 대한 정보를 나타내는 세 정수 A, B, C가 주어진다. 이는 A번 정점과 B번 정점이 가중치 C인 간선으로 연결되어 있다는 의미이다. C는 음수일 수도 있으며, 절댓값이 1,000,000을 넘지 않는다.

그래프의 정점은 1번부터 V번까지 번호가 매겨져 있고, 임의의 두 정점 사이에 경로가 있다. 최소 스패닝 트리의 가중치가 -2,147,483,648보다 크거나 같고, 2,147,483,647보다 작거나 같은 데이터만 입력으로 주어진다.

출력

첫째 줄에 최소 스패닝 트리의 가중치를 출력한다.

문제 풀이

이 문제는 이름 부터 알 수 있듯이 최소 스패닝 트리를 사용하여 문제를 풀면 됩니다!

저는 이 문제를 최소 스패닝 트리를 만드는 방법에서 크루스칼 알고리즘을 사용했습니다.

(자세한 설명은 위의 링크를 확인해주세요)

 

문제에서 입력으로 주어지는 간선들의 가중치를 기준으로 정렬한 뒤 가중치가 작은 것 부터 추가하며 사이클을 만들지 않는 트리를 만들었습니다.

 

소스 코드

import Foundation

func solution() -> Int{
    
    let firstLine = readLine()!.split(separator: " ").map({Int($0)!})
    let v = firstLine[0]
    let e = firstLine[1]
    
    var parent = Array(0...v)
    var graph: [[Int]] = []
    var lines = 0
    var result = 0
    
    // 부모 노드를 찾는 함수
    func findParent(_ index: Int) -> Int{
        
        if parent[index] == index {
            return index
        } else {
            parent[index] = findParent(parent[index])
            return parent[index]
        }
    }

    // 부모 노드를 합치는 함수
    func unionParent(_ a: Int, _ b: Int) {
        let num1 = findParent(a)
        let num2 = findParent(b)
        if a < b {
            parent[num2] = num1
        } else {
            parent[num1] = num2
        }
        
    }
    
    
    for _ in 0..<e {
        graph.append(readLine()!.split(separator: " ").map({Int($0)!}))
    }
    
    graph.sort { (a, b) -> Bool in
        return a[2] < b[2]
    }
    
    for index in 0..<graph.count {
        if lines == v - 1 {
            break
        }
        if findParent(graph[index][0]) != findParent(graph[index][1]) {
            result += graph[index][2]
            lines += 1
            unionParent(graph[index][0], graph[index][1])
        }
    }
    
    return result
}

print(solution())