프로그래머스 : 합승 택시 요금

[문제]

https://programmers.co.kr/learn/courses/30/lessons/72413

출발점과 두 사람의 집의 위치, 도로들의 예상 요금이 주어졌을 때 최저 택시요금을 구하는 문제이다.
합승은 해도 되고 안 해도 된다.

[풀이]

가중치가 모두 양인 그래프에서 최단 거리를 찾는 문제이므로 다익스트라를 이용하면 될 것 같았다.
우선 단순하게 S를 출발점으로 설정한 뒤 A, B까지의 최단거리를 구하고
A를 출발점으로 선택한 후 B까지의 최단거리를 구했다.

그리고 s_to_a + s_to_b, s_to_a + a_to_b, s_to_b + a_to_b 중 최솟값을 정답으로 리턴하게 했다.

from collections import defaultdict
import math
import heapq

def solution(n, s, a, b, fares):
    
    l = defaultdict(list)
    for start, end, fair in fares:
        heapq.heappush(l[start], [fair, end])
        heapq.heappush(l[end], [fair, start])
    
    d = [math.inf for x in range(n+1)]
    d[s] = 0
    hq = [[0, s]]
    visited = [0 for x in range(n+1)]
    
    while hq:
        dis, now = heapq.heappop(hq)
        if visited[now]: continue
        for fair, nxt in l[now]:
            if not visited[nxt]:
                d[nxt] = min(d[nxt], d[now] + fair)
                heapq.heappush(hq, [d[nxt], nxt])
        visited[now] = 1
    
    s_to_a = d[a]*1
    s_to_b = d[b]*1
    
    d = [math.inf for x in range(n+1)]
    d[a] = 0
    hq = [[0, a]]
    visited = [0 for x in range(n+1)]
    
    while hq:
        dis, now = heapq.heappop(hq)
        if visited[now]: continue
        for fair, nxt in l[now]:
            if not visited[nxt]:
                d[nxt] = min(d[nxt], d[now] + fair)
                heapq.heappush(hq, [d[nxt], nxt])
        visited[now] = 1
                
    a_to_b = d[b] * 1
    
    answer = min([s_to_a + s_to_b, s_to_a + a_to_b, s_to_b + a_to_b])
    
    return answer

하지만?
1번 테스트 케이스를 통과하지 못했다.
택시를 합승하다가 중간에 한 명이 내려 또다른 택시를 잡아 타 각자 집으로 가는 경우를 포함하지 않았기 때문이다.

이 경우에 총 요금은 임의의 노드(중간에 찢어져서 각자 가는)에서 a, b, s까지의 최저 요금의 합과 같으므로,
a, b, s를 제외한 모든 노드에서 다익스트라를 한 번씩 돌리는 과정을 추가하여 해결하였다.
노드의 갯수가 200개가 넘지 않아 이 방법을 사용해도 시간 초과가 일어나지 않는 것 같다.

    tmp = math.inf
    for i in range(1, n+1):
        if i in [a, b, s]: continue
    
        d = [math.inf for x in range(n+1)]
        d[i] = 0
        hq = [[0, i]]
        visited = [0 for x in range(n+1)]

        while hq:
            dis, now = heapq.heappop(hq)
            if visited[now]: continue
            for fair, nxt in l[now]:
                if not visited[nxt]:
                    d[nxt] = min(d[nxt], d[now] + fair)
                    heapq.heappush(hq, [d[nxt], nxt])
            visited[now] = 1

        tmp = min(tmp, d[a] + d[b] + d[s])*1
    
    answer = min([s_to_a + s_to_b, s_to_a + a_to_b, s_to_b + a_to_b, tmp])

---

다른 사람들의 풀이를 보니 플로이드 워셜 알고리즘을 쓰는 문제였나 보다.

나무위키에서 설명을 보고 문제에 맞게 구현하였따.

from collections import defaultdict
import math
import sys
import heapq

def solution(n, s, a, b, fares):
    
    d = [ [ math.inf for x in range(n) ] for x in range(n) ]
    
    for start, end, fare in fares:
        d[start-1][end-1] = fare
        d[end-1][start-1] = fare
        
    for i in range(n):
        d[i][i] = 0
        
    for m in range(n):
        for i in range(n):
            for j in range(n):
                d[i][j] = min(d[i][j], d[i][m] + d[m][j])
    
    answer = math.inf
    for i in range(n):
        answer = min(answer, d[a-1][i] + d[b-1][i] + d[s-1][i])
    
    
    return answer

두 알고리즘의 실행 시간과 메모리를 보니 플로이드 워셜은 다익스트라에 비해 저점은 높지만 고점은 낮은 모습을 보였다.