최소경로 - 전보

본 내용은 '이것이 취업을 위한 코딩테스트다 with 파이썬' 책을 기반으로 포스팅 하였습니다.

 

문제

저작권 문제가 될 수 있어 문제는 삭제합니다.

 

풀이 코드

class Heap<T: Comparable> {
    var heapArray: [[T]]
    var root: T? {
        if isMaxHeap {
            maxHeapify()
        } else {
            minHeapify()
        }
        return heapArray[0].first
    }
    var count: Int {
        return heapArray.count
    }
    var isEmpty: Bool {
        return heapArray.isEmpty
    }
    var isMaxHeap: Bool

    init(_ n: [[T]], isMaxHeap: Bool) {
        heapArray = n
        self.isMaxHeap = isMaxHeap
    }
    func isLeftChildExist(_ parent: Int) -> Bool {
        if parent * 2 + 1 >= count {
            return false
        } else {
            return true
        }
    }
    func isRightChildExist(_ parent: Int) -> Bool {
        if parent * 2 + 2 >= count {
            return false
        } else {
            return true
        }
    }
    func parentIndex(_ child: Int) -> Int {
        return child / 2
    }
    func leftChildIndex(_ parent: Int) -> Int {
        return parent * 2 + 1
    }
    func rightChildIndex(_ parent: Int) -> Int {
        return parent * 2 + 2
    }
    func maxHeapify() {
        isMaxHeap = true
        var i = count / 2
        while i >= 0 {
            var bigChild: Int = 0
            if isLeftChildExist(i) && isRightChildExist(i) {

                if heapArray[leftChildIndex(i)].first! <= heapArray[rightChildIndex(i)].first! {
                    bigChild = rightChildIndex(i)
                } else {
                    bigChild = leftChildIndex(i)
                }
            } else if isLeftChildExist(i) && !isRightChildExist(i) {
                bigChild = leftChildIndex(i)
            } else if !isLeftChildExist(i) && !isRightChildExist(i) {
                i -= 1
                continue
            }

            if heapArray[bigChild].first! >= heapArray[i].first! {
                let temp = heapArray[bigChild]
                heapArray[bigChild] = heapArray[i]
                heapArray[i] = temp
            }
            i -= 1
        }
    }
    func minHeapify() {
        isMaxHeap = false
        var i = count / 2
        while i >= 0 {
            var smallChild: Int = 0
            if isLeftChildExist(i) && isRightChildExist(i) {
                if heapArray[leftChildIndex(i)].first! >= heapArray[rightChildIndex(i)].first! {
                    smallChild = rightChildIndex(i)
                } else {
                    smallChild = leftChildIndex(i)
                }
            } else if isLeftChildExist(i) && !isRightChildExist(i) {
                smallChild = leftChildIndex(i)
            } else if !isLeftChildExist(i) && !isRightChildExist(i) {
                i -= 1
                continue
            }

            if heapArray[smallChild].first! <= heapArray[i].first! {
                let temp = heapArray[smallChild]
                heapArray[smallChild] = heapArray[i]
                heapArray[i] = temp
            }
            i -= 1
        }
    }
    func popRoot() -> [T]?{
        if isMaxHeap {
            maxHeapify()
        } else {
            minHeapify()
        }
        if !heapArray.isEmpty {
            let temp = heapArray[heapArray.count - 1]
            heapArray[heapArray.count - 1] = heapArray[0]
            heapArray[0] = temp
        }
        let rootValue = heapArray.popLast()

        return rootValue
    }
    func push(_ n: [T]) {
        heapArray.append(n)
    }
}

// -------------------------------------------------------------------------------------
// 문제 풀이 코드

// 노드의 개수, 간선의 개수, 시작노드 입력 받기
let input = readLine()!.components(separatedBy: " ").map({Int($0)!})
// 노드의 개수
let n = input[0]
// 간선의 개수
let m = input[1]
// 시작 노드
let start = input[2]
// 각 노드에 연결되어 있는 노드에 대한 정보를 저장할 배열 초기화
var graph = Array(repeating: [[Int]](), count: n+1)
// 최단 거리 정보를 저장할 배열을 모두 9999999999(무한)으로 초기화
var dp = Array(repeating: 9999999999, count: n+1)

// 모든 간선 정보 입력받기
// input[0] = 메시지 보내는 도시
// input[1] = 메시지 받을 도시 
// input[2] = 메시지를 보내는데 걸리는 시간
// 특정도시 = [걸린시간, 받을 도시]
for _ in 0..<m {
    let input = readLine()!.components(separatedBy: " ").map({Int($0)!})
    graph[input[0]].append(contentsOf: [[input[2], input[1]]])
}


func dijkstra() {

    // 큐 배열 초기화(최소 힙)
    let heap = Heap.init([[Int]](), isMaxHeap: false)

    // 시작노드로 가기위한 최단 경로를 0으로 초기화 후 큐에 삽입
    // [걸리는시간, 시작위치]
    heap.push([0, start])
    // 시작노드 최소거리 배열 초기화
    dp[start] = 0
    
    // 큐가 비어있지 않으면 반복
    while !heap.heapArray.isEmpty {
        // 큐에 저장되어 있는 정보중 최단거리가 가장 짧은 노드에 대한 정보 꺼내기 [보내는데 걸린 시간, 받을 도시]
        let dist = heap.popRoot()!
        let now = dist[1]
        
        // 현재 노드의 비용이 큐에서 꺼낸 정보의 비용보다 작다면 다음 반복 수행
        if dp[now] < dist.first! {
            continue
        }
        
        // 현재 노드와 연결되어 있는 다른 노드들을 확인
        for i in graph[now] {
            // 큐에서 꺼낸 정보의 비용 + 보내는데 걸린 비용 
            let cost = dist.first! + i[0]
            
            // 현재 노드를 거쳐서, 다른 노드로 이동하는 거리가 더 짧으면
            if cost < dp[i[1]] {
                // 현재 노드를 거쳐가는 노드에 cost 저장
                dp[i[1]] = cost
                // 큐에 [보내는데 걸린 시간, 현재 노드를 거쳐가는 노드번호] 삽입
                heap.push([cost, i[1]])
            }
        }
    }
    
    // 메시지를 받을는 도시를 카운트할 변수 선언
    var count = 0
    // 모든 메시지를 받는 데까지 걸리는 시간을 저장할 변수 선언
    var maxDist = 0
    
    for d in dp {
        if d != 9999999999 {
            count += 1
            maxDist = max(maxDist, d)
        }
    }
    print(count - 1, maxDist)
}

dijkstra()

 

문제 풀이