코딩테스트 연습 - 합승 택시 요금

플로이드 와샬 알고리즘

문제 풀이 아이디어

플로이드 와샬 알고리즘은 모든 노드에서 다른 노드까지 가는 최적의 루트를 구하는 알고리즘입니다. 이 최적의 루트를 구하는 과정에서 i에서 j로 바로 가는 루트보다 다른 k를 거쳐서 가는 것이 최적인 경우 해당 루트를 갱신합니다.

해당 문제도 같은 원리를 이용하면 됩니다. 일단 플로이드 와샬 알고리즘를 통해서 노드 간의 최단 경로를 구합니다. 노드 간의 최단 경로를 구하고 나면 출발점인 s에서 각각 a, b까지 가는 최적인 경로를 구할 차례입니다. 환승을 하는 경우가 있으므로 i라는 노드까지 같이 환승을 하게 되는 경우 비용은 (s → i) + (i → a) + (i → b)입니다.

모든 환승지점 i에 대한 완전탐색을 통해 최적의 환승 위치를 찾으면 됩니다.

코드

func solution(_ n:Int, _ s:Int, _ a:Int, _ b:Int, _ fares:[[Int]]) -> Int {
    
    // 길 없음 = 비용 무한대
    let INF = Int.max
    
    var matrix = Array(repeating: Array(repeating: INF, count: n + 1), count: n + 1)
    
    // 자기 자신으로 가는 루트 비용 0으로 초기화
    for i in 1...n {
        matrix[i][i] = 0
    }
    
    // 각 node 사이의 cost 입력
    for fare in fares {
        let c = fare[0]
        let d = fare[1]
        let f = fare[2]
        matrix[c][d] = f
        matrix[d][c] = f
    }
    
    // 플로이드 와샬 알고리즘을 통해서 각 node 사이의 최적루트 구하기
    for k in 1...n {
        for i in 1...n {
            for j in 1...n {
                // arithmetic overflow를 방지하기 위한 guard문
                guard matrix[i][k] < INF && matrix[j][k] < INF else { continue }
                matrix[i][j] = min(matrix[i][j], matrix[i][k] + matrix[k][j])
            }
        }
    }
    
    // 환승을 하지 않는 경우와 환승을 하는 경우의 비용 비교
    var ans = matrix[s][a] + matrix[s][b]
    
    // i를 거쳐 환승하는 경우의 cost를 구하고 ans와 비교
    for i in 1...n {
        guard matrix[s][i] < INF && matrix[i][a] < INF && matrix[i][b] < INF else { continue }
        ans = min(ans, matrix[s][i] + matrix[i][a] + matrix[i][b])
    }

    return ans
}

다익스트라 알고리즘

다익스트라 알고리즘이 더 빠를까?

일반적인 경우

구현의 난이도는 위에서 사용한 방법인 플로이드-와샬 알고리즘이 더 낮습니다. 하지만 플로이드-와샬 알고리즘의 경우 모든 노드에서 모든 노드 사이의 최단 거리를 구하기 때문에 시간복잡도가 높습니다. 안에 3중 반복문이 있으므로 시간복잡도는 노드의 갯수를 N이라고 했을 때 O(N^3)입니다.

하지만 우리가 필요한 최적의 경로는 출발점에서 목적지, 출발점에서 합승지점, 합승지점에서 목적지까지의 경로입니다. 즉 모든 경로를 구할 필요가 없다는 것입니다.

따라서 한 노드에서 모든 노드까지의 최단거리를 구하는 길찾기 알고리즘인 다익스트라 알고리즘을 사용해보도록 하겠습니다. 다익스트라 알고리즘의 시간복잡도는 간선의 갯수를 E라고 했을 때 O(ElogE)입니다. 이 문제에서 간선의 갯수는 fares 배열의 크기입니다. 최대 N(N-1)/2라고 합니다. Big-O 기준으로 상수항을 모두 제거하면 N^2이 되겠네요. 그렇다면 N으로 표현한 시간복잡도는 O(N^2 logN^2) = O(N^2 2logN) = O(N^2 * logN)입니다. ✅  즉 일반적으로는 플로이드-와샬의 O(N^3)보다 효율적인 알고리즘이 되는 것이죠.

이 문제에서는?

하지만 여기서 문제가 하나 있습니다. 다익스트라 알고리즘을 1번만 실행할 경우는 물론 플로이드-와샬보다 더 빠른 알고리즘입니다만 우리는 모든 환승지에 대해서 다익스트라 알고리즘을 실행해야 합니다.

즉 1번만 실행하는 것이 아니라 아래의 반복문 안에서 다익스트라 알고리즘을 실행해야 한다는 것입니다.

// 다익스트라(출발, 도착)

// 모든 환승지에 대해서 다익스트라 실시
for i in 1...n {
	환승 = i
	ans = min(ans, 다익스트라(출발, 환승) + 다익스트라(환승, a집) + 다익스트라(환승, b집)
}

위와 같은 반복문 안에서 다익스트라를 실시하는 것까지 감안하면 시간복잡도는 O(N^2 logN 3N) = O(N^3 * logN)이 되는 것입니다. 오히려 다익스트라 알고리즘으로 푸는 것이 더 시간이 걸리는 셈입니다.

실제로 두 풀이를 비교해보면 실행시간의 차이가 있는 것을 볼 수 있습니다.

문제 풀이 아이디어

더 간단한 구현에 더 빠른 속도인 플로이드-와샬 알고리즘을 사용하지 않을 이유가 없지만 경험 삼아 다익스트라 알고리즘을 사용해서 풀어보도록 하겠습니다. 개인적으로 다익스트라 알고리즘을 복습할 필요가 있기도 하고요.

일단 다익스트라 알고리즘을 구현합니다. 구현 방법은 별도의 포스팅에 정리하도록 하겠습니다.

다익스트라 알고리즘을 구현한 함수를 통해서 우리는 해당 트리에서 한 노드에서 다른 노드까지 가는 최단거리를 구할 수 있습니다. 이 함수를 활용해서 일단 환승을 하지 않고 각각 따로 간 경우의 비용, 즉 dijstra(s, a) + dijstra(s, b)를 구합니다.

그리고 모든 환승지점 i에 대해서 dijstra(s, i) + dijstra(i, a) + dijstra(i, b)를 비교하면서 최소값을 찾으면 됩니다.

코드

import Foundation

// Swift로 최소힙 구현
struct MinHeap<T: Comparable> {
    var heap: [T] = []

    var isEmpty: Bool {
        return heap.count <= 1 ? true : false
    }
    
    
    init() {}

    init(_ element: T) {
        heap.append(element)
        heap.append(element)
    }
    
    mutating func insert(_ element: T) {
        
        if heap.isEmpty {
            heap.append(element)
            heap.append(element)
            return
        }
        
        heap.append(element)
        
        func isMoveUp(_ insertIndex: Int) -> Bool {
            if insertIndex <= 1 {
                return false
            }
            
            let parentIndex = insertIndex / 2
            return heap[insertIndex] < heap[parentIndex] ? true : false
        }
        
        var insertIndex = heap.count - 1
        
        while isMoveUp(insertIndex) {
            let parentIndex = insertIndex / 2
            heap.swapAt(insertIndex, parentIndex)
            insertIndex = parentIndex
        }
    }
    

    enum moveDownDirection { case left, right, none }
    

    mutating func pop() -> T? {
        
        if heap.count <= 1 {
            return nil
        }
        
        let returnData = heap[1]
        heap.swapAt(1, heap.count - 1)
        heap.removeLast()
        
        func moveDown(_ poppedIndex: Int) -> moveDownDirection {
            
            let leftChildIndex = poppedIndex * 2
            let rightChildIndex = leftChildIndex + 1
            
            if leftChildIndex >= heap.count {
                return .none
            }
            
            if rightChildIndex >= heap.count {
                return heap[leftChildIndex] < heap[poppedIndex] ? .left : .none
            }
            
            if (heap[leftChildIndex] > heap[poppedIndex]) && (heap[rightChildIndex] > heap[poppedIndex]) {
                return .none
            }
            
            if (heap[leftChildIndex] < heap[poppedIndex]) && (heap[rightChildIndex] < heap[poppedIndex]) {
                return heap[leftChildIndex] < heap[rightChildIndex] ? .left : .right
            }
            
            if (heap[leftChildIndex] < heap[poppedIndex]) || (heap[rightChildIndex] < heap[poppedIndex]) {
                return heap[leftChildIndex] < heap[rightChildIndex] ? .left : .right
            }
            
            return .none
        }
        
        var poppedIndex = 1
        while true {
            switch moveDown(poppedIndex) {
            case .none:
                return returnData
            case .left:
                let leftChildIndex = poppedIndex * 2
                heap.swapAt(poppedIndex, leftChildIndex)
                poppedIndex = leftChildIndex
            case .right:
                let rightChildIndex = (poppedIndex * 2) + 1
                heap.swapAt(poppedIndex, rightChildIndex)
                poppedIndex = rightChildIndex
            }
        }
    }
}

// Node 구조체 구현
    // 거리 비교를 위해서 Comparable
struct Node: Comparable {
    
    let node: Int
    let distance: Int
    
    static func < (lhs: Node, rhs: Node) -> Bool {
        return lhs.distance < rhs.distance
    }
    
}

func solution(_ n:Int, _ s:Int, _ a:Int, _ b:Int, _ fares:[[Int]]) -> Int {
    
    // 다익스트라 알고리즘 (from -> to의 최소 비용)
    func dijkstra(_ from: Int, _ to: Int) -> Int {
        
        var distances = Array(repeating: INF, count: n + 1)
        var heap = MinHeap<Node>()
        
        heap.insert(Node(node: from, distance: 0))
        distances[from] = 0
        
        while !heap.isEmpty {
            guard let popped = heap.pop() else { break }
            let now = popped.node
            let dist = popped.distance
            
            if distances[now] < dist { continue }
            
            for i in 1...n {
                if matrix[now][i] < INF {
                    let cost = dist + matrix[now][i]
                    if cost < distances[i] {
                        distances[i] = cost
                        heap.insert(Node(node: i, distance: cost))
                    }
                }
            }
        }
        
        return distances[to]
    }

    // 길 없음 = 비용 무한대
    let INF = Int.max

    var matrix = Array(repeating: Array(repeating: INF, count: n + 1), count: n + 1)

    // 자기 자신으로 가는 루트 비용 0으로 초기화
    for i in 1...n {
        matrix[i][i] = 0
    }

    // 각 node 사이의 cost 입력
    for fare in fares {
        let c = fare[0]
        let d = fare[1]
        let f = fare[2]
        matrix[c][d] = f
        matrix[d][c] = f
    }

    // 환승을 하지 않는 경우와 환승을 하는 경우의 비용 비교
    var ans = dijkstra(s, a) + dijkstra(s, b)

    // i를 거쳐 환승하는 경우의 cost를 구하고 ans와 비교
    for i in 1...n {
        let costAB = dijkstra(s, i)
        let costA = dijkstra(i, a)
        let costB = dijkstra(i, b)
        guard costAB < INF && costA < INF && costB < INF else { continue }
        ans = min(ans, costAB + costA + costB)
    }

    return ans
}