[Swift] BOJ 1753 최단경로

BOJ 1753 최단경로

 

접근 방식

 

다익스트라의 대표 문제. 다만 스스로 통과시키지는 못함.

 

다익스트라 로직 자체는 맞았는데 계속 시간 초과가 났음. 원인은 우선순위 큐를 힙이 아니라 배열로 두고 매 반복마다 sorted로 다시 정렬한 부분이었음(아래 시행착오 코드). 정렬 비용이 반복마다 붙어 정점과 간선이 많아지면 감당이 안 됨.

 

힙 기반 우선순위 큐로 바꿔야 풀리는 구조라, 힙을 구현한 다른 분의 코드를 학습함. 정답 코드는 이진 힙으로 push와 pop을 O(log n)에 처리함. 힙부터 제대로 공부하고 다시 돌아올 계획.

 

시행착오

 

내 코드 - 시간초과

 

struct b1753 {
    static func run() {
        // 입력부
        let firstLine = readLine()!.split(separator: " ").map( {Int(String($0))! })

        let secondLine = Int(readLine()!)!
        
        var table = [Int](repeating: Int(1e9), count: firstLine[0]+1)
        var graph = [Int: [Int: Int]]()
        
        for i in 1...firstLine[0] {
            graph[i] = [:]
        }
        
        for _ in 0..<firstLine[1] {
            let info = readLine()!.split(separator: " ").map { Int($0)! }
            
            let start = info[0]
            let end = info[1]
            let weight = info[2]
            
            var temp = graph[start]!
            
            if temp[end] == nil {
                temp[end] = weight
            } else {
                temp[end] = min(temp[end]!, weight)
            }
            graph[start] = temp
        }
        
        // 구현부 - 다익스트라 알고리즘을 구현합니다.
        
        var priorityQueue = [(Int, Int)]()
        
        var visited = [secondLine] // 방문한 노드
        table[secondLine] = 0
        
        graph[secondLine]!.forEach {
            
            priorityQueue.append(($0.key,$0.value))
        }
        
        priorityQueue = priorityQueue.sorted { $0.1 > $1.1 }
        
        var popNode = secondLine
        while !priorityQueue.isEmpty {
            
            // 테이블 업데이트
            priorityQueue.forEach { (node, weight) in
                if !visited.contains(node) {
                    table[node] = min(table[node], weight + table[popNode])
                }
                
            }
            
            // 우선순위 큐 업데이트 -> 다음 작업을 위함.
            popNode = priorityQueue.last!.0
            visited.append(popNode)
            priorityQueue.removeAll()
            
            graph[popNode]?.forEach{
                priorityQueue.append(($0.key,$0.value))
            }
            priorityQueue = priorityQueue.sorted { $0.1 > $1.1 }
            
        }
        
        // 출력부
        table[1...].forEach {
            if $0 == Int(1e9) {
                print("INF")
            } else {
                print($0)
            }
        }
    }
}

 

코드

 

정답 코드이며, 출처는 제일 아래에 있음.

 

import Foundation

struct PriorityQueue<T> {
  var array = [T]()
  let sort: (T, T) -> Bool

  init(sort: @escaping (T, T) -> Bool) {
    self.sort = sort
  }

  var isEmpty: Bool {
    return array.isEmpty
  }

  var count: Int {
    return array.count
  }

  func peek() -> T? {
    return array.first
  }

  func leftChildIndex(ofParentAt index: Int) -> Int {
    return (2 * index) + 1
  }

  func rightChildIndex(ofParentAt index: Int) -> Int {
    return (2 * index) + 2
  }

  func parentIndex(ofChildAt index: Int) -> Int {
    return (index - 1) / 2
  }

  // MARK:- remove operation
  mutating func pop() -> T? {
    guard !isEmpty else {
      return nil
    }

    array.swapAt(0, count - 1)
    defer {
      siftDown(from: 0)
    }
    return array.removeLast()
  }

  mutating func siftDown(from index: Int) {
    var parent = index
    while true {
      let left = leftChildIndex(ofParentAt: parent)
      let right = rightChildIndex(ofParentAt: parent)
      var candidate = parent

      if left < count && sort(array[left], array[candidate]) {
        candidate = left
      }
      if right < count && sort(array[right], array[candidate]) {
        candidate = right
      }
      if candidate == parent {
        return
      }
      array.swapAt(parent, candidate)
      parent = candidate
    }
  }

  // MARK:- insert operation
  mutating func push(_ element: T) {
    array.append(element)
    siftUp(from: array.count - 1)
  }

  mutating func siftUp(from index: Int) {
    var child = index
    var parent = parentIndex(ofChildAt: child)
    while child > 0 && sort(array[child], array[parent]) {
      array.swapAt(child, parent)
      child = parent
      parent = parentIndex(ofChildAt: child)
    }
  }
}

let input = readLine()!.split(separator: " ").map { Int(String($0))! }
let k = Int(readLine()!)!
let v = input[0], e = input[1]
let INF = Int(1e9)

var edges = [[(Int, Int)]](repeating: [(Int, Int)](), count: v + 1)

for _ in 0..<e {
  let i = readLine()!.split(separator: " ").map { Int(String($0))! }
  edges[i[0]].append((i[1], i[2]))
}

var dist = [Int](repeating: INF, count: v + 1)
var pq = PriorityQueue<(Int, Int)>(sort: { $0.1 < $1.1 })

dist[k] = 0
pq.push((k, 0))

while !pq.isEmpty {
  let (now, cost) = pq.pop()!

  if cost > dist[now] {
    continue
  }

  for (next, w) in edges[now] {
    if dist[now] + w <= dist[next] {
      dist[next] = dist[now] + w
      pq.push((next, dist[next]))
    }
  }
}

for i in 1...v {
  if dist[i] < INF {
    print(dist[i])
  } else {
    print("INF")
  }
}

 

참고

 

https://github.com/skyqnaqna/algorithm_study/blob/main/boj/Dijkstra/BOJ_1753.swift#L8

 

GitHub - skyqnaqna/algorithm_study: 알고리즘 문제 코드 기록 알고리즘 문제 코드 기록. Contribute to skyqnaqna/algorithm_study development by creating an account on GitHub.