[백준/Swift] 1238 : 파티/ 우선순위 큐 개념 정리, 최소 힙 개념 with 라이노님 Heap 정리

문제

1238번: 파티

간단한 문제 요약

 각 마을에는 한명의 학생이 살고 있습니다. n명의 학생은 x번 마을에 모여 파티를 합니다. 이때 m개의 도로는 단 방향입니다. i 비용을 들여 도로를 지나갑니다.

 

고려해야 할 사항

• N, M, X 입력은 1 부터입니다.. (X제약을 못봐서 계속 X위치를 다르게 구했다는..ㅠㅠ)

문제 풀이, 했갈렸던 점

 x번을 제외한 마을 학생들은 각자가 갈 수있는 모든 마을에 대해 최단 경로를 구한 후 x일때의 최단 경로를 저장합니다. 이후 x번에서 갈 수있는 최단 경로를 추가로 구한 후 이전에 구했던 최단 경로와 더한 후 max값을 찾으면 됩니다. 한 가지 했갈렸던 점은 x번의 마을에 있는 학생은 어떻게 해야하는지에 대한 고민이 있었는데 x번 학생은 x번 마을에 그대로 있는 것이었습니다.

코드

// 힙 구조는 아래에 있습니다.

struct Node {
    let vertex: Int
    let cost: Int
}
extension Node: Equatable, Comparable {
    static func <(lhs: Node, rhs: Node) -> Bool {
        return lhs.cost < rhs.cost
    }
}

//MARK: - Properties
var input = readLine()!.split{$0==" "}.map{Int(String($0))!}
let (n,m,x) = (input[0],input[1],input[2] - 1)
var graph = Array(repeating: [Node](), count: n)

var res = Array(repeating: 0, count: n)
var lastDist = res
(0..<m).forEach { _ in
    let temp = readLine()!.split{$0==" "}.map{Int(String($0))!}
    let (s,e,c) = (temp[0]-1,temp[1]-1,temp[2])
    graph[s].append(Node(vertex: e, cost: c))
}

func dijkstra(_ start: Int, _ last: Bool = false) {
    var dist = Array(repeating: Int.max, count: n)
    dist[start] = 0
    var priority = Heap<Node>()
    priority.insert(Node(vertex: start, cost: 0))
    while let curNode = priority.delete() {
        if dist[curNode.vertex] < curNode.cost { continue }
        for node in graph[curNode.vertex] {
            if dist[curNode.vertex] + node.cost < dist[node.vertex] {
                dist[node.vertex] = dist[curNode.vertex] + node.cost
                priority.insert(Node(vertex: node.vertex, cost: dist[node.vertex]))
            }
        }
    }
    guard last else {
        res[start] = dist[x]
        return
    }
    lastDist = dist
}

func solution() {
    (0..<n).forEach { dijkstra($0) }
    dijkstra(x,true)
    (0..<res.count).forEach { res[$0] += lastDist[$0] }
    print(res.max()!)
}

solution()

라이노님 힙 구조 조금 변형

struct Heap<T> where T: Comparable {
    private var heap = [T]()
    let compare: (T,T) -> Bool
    init(compare: @escaping (T,T) -> Bool) {
        self.compare = compare
    }
    init() {
        self.init(compare: <)
    }
}
extension Heap {
    var isEmpty: Bool {
        return heap.isEmpty
    }
    mutating func insert(_ element: T) {
        var idx = heap.count
        heap.append(element)
        while idx > 0 && compare(heap[idx], heap[(idx-1)/2]) {
            heap.swapAt((idx-1)/2, idx)
            idx = (idx-1)/2
        }
    }
    mutating func delete() -> T? {
        guard !heap.isEmpty else {
            return nil
        }
        guard heap.count != 1 else {
            return heap.removeFirst()
        }
        let res = heap.first!
        heap.swapAt(0, heap.count-1)
        _ = heap.removeLast()
        var idx = 0
        while idx < heap.count {
            let (left,right) = (idx*2+1,idx*2+2)
            if right < heap.count {
                if compare(heap[left], heap[right]) && compare(heap[left], heap[idx]) {
                    heap.swapAt(left, idx)
                    idx = left
                }else if compare(heap[right],heap[idx]) {
                    heap.swapAt(right, idx)
                    idx = right
                }else {
                    break
                }
            }else if left < heap.count {
                if compare(heap[left], heap[idx]) {
                    heap.swapAt(left, idx)
                    idx = left
                }else {
                    break
                }
            }else {
                break
            }
        }
        return res
    }
}

라이노님 github