Algorithm/Programmers

Programmers) Lv3 섬 연결하기

삼쓰 웅쓰 2020. 3. 6. 19:39
728x90

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

분류: Lv3 , Greedy


MST의 대표적인 문제입니다!!!

Prim 알고리즘으로 해결해봤는데요. 자세한 설명은 이론 설명 그 자체일 듯 합니당
다음글을 참고해주세용 :)

/**
 * 현재 노드에 연결되어 있는 노드만 확인할 것이기 때문에
 * 간단하게 목적지와 비용만 저장하는 Edge 를 정의합니다.
 */
struct Edge {
    var destination: Int
    var cost: Int
}

func updateDist(_ tree: [Edge], dist: inout [Int?]) -> [Edge] {
    for edge in tree {
        if let d = dist[edge.destination] {
            if d > edge.cost {
                dist[edge.destination] = edge.cost
            }
        } else {
            dist[edge.destination] = edge.cost
        }
    }
    
    return tree
}

func solution(_ n: Int, _ costs: [[Int]]) -> Int {
    
    // 최소 비용을 저장할 배열
    var dist = [Int?](repeating: nil, count: n)
    
    // 방문 체크 배열
    var visit = [Bool](repeating: false, count: n)
    
    // 방문한 노드의 개수 ( 방문 체크 배열에서 확인해도 되지만.. 그냥 간단하게 변수를 하나 만들었습니다. )
    var visitCount = 0
    
    // 현재 방문 노드를 저장합니다. 0에서 시작합니다.
    var now = 0
    
    // 최소비용을 저장한다.
    var minimumCost = 0
    
    /**
     * 배열의 인덱스가 해당 노드가 되고 연결된
     * 연결된 선을 저장합니다.
     */
    var tree = [[Edge]](repeating: [], count: n)
    
    /**
     * Setup Tree
     * 선의 양 끝을 둘다 저장해줘야 합니다.
     */
    for cost in costs {
        tree[cost[0]].append(Edge(destination: cost[1], cost: cost[2]))
        tree[cost[1]].append(Edge(destination: cost[0], cost: cost[2]))
    }
    
    visit[now] = true
    visitCount += 1
    
    // 탐색
    while visitCount < n {
        let tree = updateDist(tree[now], dist: &dist)
        
        // 최소값 구하기
        var minCost: Int? = nil
        for i in 0..<n {
            guard !visit[i], let d = dist[i] else {continue}
            
            if let min = minCost {
                if min > d {
                    minCost = d
                    now = i
                }
            } else {
                minCost = d
                now = i
            }
        }
        
        visit[now] = true
        visitCount += 1
        minimumCost += minCost!
    }
    
    
    return minimumCost
}