[Swift/프로그래머스] 섬 연결하기 ( 그리디, MST )

https://school.programmers.co.kr/learn/courses/30/lessons/42861

프로그래머스

 

바로 이전에 포스팅한 MST를 활용해 해결하는 문제다.

 

가중치가 나오고 그래프 문제라고 한다면 무조건 MST로 풀면 된다고 생각하자.

 

준비물

1. 부모배열 ( 자기 자신을 담은 )

2. find ( 부모찾기 함수 , 갱신도 해줘야함 )

3. union ( 서로 다른 부모를 같은 부모로 갱신 )

4. 가중치를 기준으로 오름차순 정렬

5. 순환이 없어야함 ( union 함수로 해결 )

 

풀이

import Foundation

func solution(_ n: Int, _ costs: [[Int]]) -> Int {
    
    // 가중치를 기준으로 오름차순 정렬
    let sortedLand = costs.sorted{ $0[2] < $1[2] }
    
    // 부모 배열 생성
    var parent = [Int]()
    
    for i in 0..<n{
        parent.append(i)
    }
    
    func find(_ parent: inout [Int], _ index: Int) -> Int { // 부모번호 반환
        if parent[index] == index {
            return index
        }
        
        parent[index] = find(&parent, parent[index])
        return parent[index]
    }
    
    func union(_ parent : inout [Int], _ a: Int, _ b: Int) {
        let parentOfA = find(&parent, a)
        let parentOfB = find(&parent, b)
        
        if parentOfA < parentOfB {
            parent[parentOfB] = parentOfA
        }else{
            parent[parentOfA] = parentOfB
        }
    }
    
    var result = 0
    
    for cost in sortedLand {
        let a = cost[0]
        let b = cost[1]
        let price = cost[2]
        if find(&parent,a) != find(&parent,b){
            union(&parent,a,b)
            result += price
        }
    }
    
    return result
}