프로그래머스 - 가장 먼 노드(Lv. 3)

문제 링크

풀이

import Foundation

func solution(_ n:Int, _ edge:[[Int]]) -> Int {
    var graph : [Int: [Int]] = [:]
    // 일단 간선 정보는 뺀 그래프 생성
    for index in 1...n {
        graph[index] = []
    }
    
    // 간선 정보 넣기
    for pathInfo in edge {
        graph[pathInfo[0]]!.append(pathInfo[1])
        graph[pathInfo[1]]!.append(pathInfo[0])
    }
    
    // 다익스트라 (각 노드 최단거리 찾기)
    let shortestInfo = dijkstra(graph, 1)
    
    // 제일 길었던 노드 개수 세기
    var farDist = 0
    var farNodeCount = 0
    for key in shortestInfo.keys {
        let dist = shortestInfo[key]!
        if dist > farDist {
            farDist = dist
            farNodeCount = 1
        } else if dist == farDist {
            farNodeCount += 1
        }
    }
    return farNodeCount
}

// ex) graph - [1: [2, 3], 2: [3, 4, 5], ...]
// 1에서는 2, 3을 갈 수 있고 2에서는 3, 4, 5 노드로 갈 수 있다는 내용
func dijkstra(_ graph : [Int: [Int]], _ start : Int) -> [Int: Int] {
    var distances: [Int:Int] = [:] // 최단거리 정보
    for item in graph {
        distances.updateValue(Int.max, forKey: item.key)
    }
    // 본인 출발 도착 0
    distances[start] = 0
    var queue = PriorityQueue()
    queue.enqueue(Node(start, 0))
    
    while queue.count != 0 {
        let dequeued = queue.dequeue()
        let currentNode = dequeued.node
        let currentDistance = dequeued.dist
        
        if distances[currentNode]! < currentDistance {
            continue
        }
        
        for node in graph[currentNode]! {
            let distance = currentDistance + 1
            if distance < distances[node]! {
                distances[node] = distance
                queue.enqueue(Node(node, distance))
            }
        }
    }
    
    return distances
}

struct Node: Equatable  {    
    let node: Int   // 노드 번호
    let dist: Int   // 지나온 거리의 합
    
    init(_ node: Int, _ dist: Int) {
        self.node = node
        self.dist = dist
    }
    
    static func >(lhs: Node, rhs: Node) -> Bool {
        return lhs.dist > rhs.dist
    }
    
    static func <(lhs: Node, rhs: Node) -> Bool {
        return lhs.dist < rhs.dist
    }
}

struct PriorityQueue {
    private var nodes: [Node] = []

    var count: Int {
        return nodes.count
    }

    mutating func enqueue(_ element: Node) {
        nodes.append(element)
        let lastIndex = count - 1
        shiftUp(child: lastIndex)
    }

    mutating func dequeue() -> Node {        
        nodes.swapAt(0, count - 1)
        defer { shiftDown(parent: 0) }
        return nodes.removeLast()
    }
    
    private func parentIndex(of child: Int) -> Int {
        return (child - 1) / 2
    }
    
    private func leftChildIndex(of parent: Int) -> Int {
        return (parent * 2) + 1
    }
    
    private func rightChildIndex(of parent: Int) -> Int {
        return parent * 2
    }
    
    mutating private func shiftUp(child: Int) {
        var child = child
        var parent = parentIndex(of: child)
        
        while child > 0 && nodes[child] < nodes[parent] {
            nodes.swapAt(child, parent)
            child = parent
            parent = parentIndex(of: child)
        }
    }
    
    mutating private func shiftDown(parent: Int) {
        var parent = parent
        
        while true {
            let left = leftChildIndex(of: parent)
            let right = rightChildIndex(of: parent)
            var candidate = parent
            
            if right < nodes.count && nodes[right] < nodes[candidate] {
                candidate = right
            } else if left < nodes.count && nodes[left] < nodes[candidate] {
                candidate = left
            } else {
                return
            }
            
            nodes.swapAt(parent, candidate)
            parent = candidate
        }
    }
}

후기

'가장 먼 노드를 찾으시오' 라는 지문을 읽자마자 바로 다익스트라 알고리즘을 사용할 때구나라는게 번뜩였다.
처음 다익스트라 알고리즘을 구현할 때 우선순위 큐를 사용하지 않은 간소화(야매) 방식을 사용하여 풀었는데 테스트 케이스의 반이 시간초과로 실패되었다.
우선순위 큐를 도입하자 다 통과되는 걸 보고 알고리즘의 대단함을 다시 느낀 문제

(Swift는 왜 큐가 기본적으로 없는지 참...)