[백준 2892] 국왕의 방문

Junyoung Park·2022년 8월 5일
0

코딩테스트

목록 보기
525/631
post-thumbnail

1. 문제 설명

국왕의 방문

2. 문제 분석

다익스트라 알고리즘을 구현할 때 현재 시점의 국왕이 어디에 있는지를 확인해야 한다. intersectionData를 통해 특정 시간을 입력값으로 넣을 때 현 노드 ~ 다음 노드, 머무르는 시간을 리턴하는 함수를 구현했다. 이를 통해 현 시점의 다익스트라 알고리즘이 맞닥뜨린 국왕을 기다리는 waitCost를 구현할 수 있다. waitCost를 일반 다익스트라 알고리즘의 nextCost처럼 관리한다면 국왕을 맞닥뜨리는 케이스와 일반적인 케이스 모두를 다룰 수 있다.

3. 나의 풀이

import Foundation

struct PQ<T> {
    var nodes = [T]()
    let sort:(T, T) -> Bool
    init(sort:@escaping(T, T) -> Bool) {
        self.sort = sort
    }
    var isEmpty: Bool {
        return nodes.isEmpty
    }
    var count: Int {
        return nodes.count
    }
    func leftChild(from parentIndex: Int) -> Int {
        return parentIndex * 2 + 1
    }
    func rightChild(from parentIndex: Int) -> Int {
        return parentIndex * 2 + 2
    }
    func parentIndex(from child: Int) -> Int {
        return (child - 1) / 2
    }
    mutating func shiftDown(from index: Int) {
        var parent = index
        while true {
            let left = leftChild(from: parent)
            let right = rightChild(from: parent)
            var candidate = parent
            if left < count && sort(nodes[left], nodes[candidate]) {
                candidate = left
            }
            if right < count && sort(nodes[right], nodes[candidate]) {
                candidate = right
            }
            if candidate == parent {
                return
            }
            nodes.swapAt(parent, candidate)
            parent = candidate
        }
    }
    mutating func shiftUp(from index: Int) {
        var child = index
        var parent = parentIndex(from: child)
        while child > 0 && sort(nodes[child], nodes[parent]) {
            nodes.swapAt(child, parent)
            child = parent
            parent = parentIndex(from: child)
        }
    }
    mutating func pop() -> T? {
        guard !isEmpty else { return nil}
        nodes.swapAt(0, count-1)
        defer {
            shiftDown(from: 0)
        }
        return nodes.removeLast()
    }
    mutating func push(_ element: T) {
        nodes.append(element)
        shiftUp(from: count-1)
    }
}

let INF = Int.max
var input = readLine()!.split(separator: " ").map{Int(String($0))!}
let (N, M) = (input[0], input[1])
var nodes = Array(repeating: [(Int, Int)](), count: N+1)
input = readLine()!.split(separator: " ").map{Int(String($0))!}
let (A, B, K, G) = (input[0], input[1], input[2], input[3])
let intersections = Array(readLine()!.split(separator: " ").map{Int(String($0))!})
var intersectionDict = [Int:Int]()
var intersectionIdx = 0
for intersection in intersections {
    intersectionDict[intersection] = intersectionIdx
    intersectionIdx += 1
}
var intersectionTimes = Array(repeating: 0, count: G-1)
for _ in 0..<M {
    let input = readLine()!.split(separator: " ").map{Int(String($0))!}
    let (node1, node2, cost) = (input[0], input[1], input[2])
    nodes[node1].append((node2, cost))
    nodes[node2].append((node1, cost))
    if let idx1 = intersectionDict[node1], let idx2 = intersectionDict[node2], abs(idx1-idx2) == 1 {
        if idx1 < idx2 {
            intersectionTimes[idx1] = cost
        } else {
            intersectionTimes[idx2] = cost
        }
    }
}

for i in 1..<G-1 {
    intersectionTimes[i] += intersectionTimes[i-1]
}

let distances = Dijkstra(start: A, K: K)
let answer = distances[B] - K
print(answer)

func getIntersectionData(curCost: Int) -> (Int, Int, Int) {
    for i in 0..<G-1 {
        if curCost < intersectionTimes[i] {
            return (intersections[i], intersections[i+1], intersectionTimes[i])
        }
    }
    return (-1, -1, -1)
}

func Dijkstra(start: Int, K: Int) -> [Int] {
    var distances = Array(repeating: INF, count: N+1)
    distances[start] = K
    var pq = PQ<(Int, Int)>(sort: {$0.0 < $1.0})
    pq.push((K, start))
    
    while !pq.isEmpty {
        guard let curData = pq.pop() else { break }
        let curCost = curData.0
        let curNode = curData.1
        
        if distances[curNode] < curCost {
            continue
        }
        
        let intersectionData = getIntersectionData(curCost: curCost)
        let interNode1 = intersectionData.0
        let interNode2 = intersectionData.1
        let interCost = intersectionData.2
        
        for nextData in nodes[curNode] {
            let nextNode = nextData.0
            let nextCost = nextData.1
            
            var waitCost = 0
            if (interNode1 == curNode && interNode2 == nextNode) || (interNode1 == nextNode && interNode2 == curNode) {
                waitCost = interCost - curCost
            }
            let totalCost = waitCost + curCost + nextCost
            
            if distances[nextNode] > totalCost {
                distances[nextNode] = totalCost
                pq.push((totalCost, nextNode))
            }
        }
    }
    return distances
}
profile
JUST DO IT

0개의 댓글