문제

https://www.acmicpc.net/problem/1753

풀이

다익스트라 입문 문제이다. 단, 스위프트로 풀려면 힙을 직접 구현해야 한다.

입력으로 주어지는 엣지의 웨이트가 모두 같다는 조건이 없으므로, BFS가 아닌 다익스트라로 탐색해야 한다.

코드

import Foundation

struct Heap {
    var heap: [(Int, Int)] = []
    
    func isEmpty() -> Bool {
        return heap.isEmpty ? true : false
    }

    mutating func insert(_ value: (Int, Int)) {
        heap.append(value)
        var currentIndex = heap.count - 1
        while currentIndex > 0 {
            let parentIndex = (currentIndex - 1) / 2
            if heap[currentIndex].1 < heap[parentIndex].1 {
                heap.swapAt(currentIndex, parentIndex)
                currentIndex = parentIndex
            } else {
                break
            }
        }
    }

    mutating func deleteMin() -> (Int, Int) {
        if heap.isEmpty{
            return (0, 0)
        }
        let min = heap[0]
        heap[0] = heap[heap.count - 1]
        heap.removeLast()
        
        var currentIndex = 0
        while true {
            let leftChildIndex = 2 * currentIndex + 1
            let rightChildIndex = 2 * currentIndex + 2
            
            if leftChildIndex >= heap.count {
                break
            }
            var minChildIndex = leftChildIndex
            if rightChildIndex < heap.count && heap[rightChildIndex].1 < heap[leftChildIndex].1 {
                minChildIndex = rightChildIndex
            }
            if heap[minChildIndex].1 < heap[currentIndex].1 {
                heap.swapAt(currentIndex, minChildIndex)
                currentIndex = minChildIndex
            } else {
                break
            }
        }
        return min
    }
}

func dijkstra(k: Int) {
    var heap = Heap()
    heap.insert((k, 0))
    distanceTable[k] = 0
    while !heap.isEmpty() {
        let edge = heap.deleteMin()
        if distanceTable[edge.0] < edge.1 {
            continue
        }
        for node in graph[edge.0] {
            let cost = edge.1 + node.1
            if cost < distanceTable[node.0] {
                distanceTable[node.0] = cost
                heap.insert((node.0, cost))
            }
        }
    }
}

let ve = readLine()!.split(separator: " ").map { Int(String($0))! }
let k = Int(readLine()!)!
var graph = [[(Int, Int)]](repeating: [], count: ve[0] + 1)
for _ in 0..<ve[1] {
    let uvw = readLine()!.split(separator: " ").map { Int(String($0))! }
    graph[uvw[0]].append((uvw[1], uvw[2]))
}
var distanceTable = [Int](repeating: 300_001, count: ve[0] + 1)

dijkstra(k: k)

for i in 1...ve[0] {
    if distanceTable[i] == 300_001 {
        print("INF")
    } else {
        print(distanceTable[i])
    }
}