[프로그래머스-파이썬] 합승 택시 요금

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

프로그래머스

 

크게 두 가지 경우의 수가 있다.

1. s->a, s->b로 바로 가는 방법

2. s->k(중간지점)->a, s->k->b 중간까지 같이 타고 가다가 각자 타고 가는 방법

- 양방향 그래프이므로 k-> a를 구하는 것은 a->k를 구하는 것과 같다.

따라서 먼저 s, a, b을 각 출발점으로 두고 각 노드를 방문하는 코드를 짜서 이를 비교하면 된다.

 

import heapq

def solution(n, s, a, b, fares):  
    answer = float('inf')
    graph = [[] for _ in range(n+1)]
    
    for f in fares:
        graph[f[0]].append([f[1], f[2]])
        graph[f[1]].append([f[0], f[2]])
    
    def search(cost, start):
        # nonlocal graph
        cost[start] = 0
        q = []
        heapq.heappush(q, (start, 0))

        while q:
            cur, cur_cost = heapq.heappop(q)
            if cost[cur] < cur_cost:
                continue

            for node, dis in (graph[cur]):
                if cur_cost + dis < cost[node]:
                    cost[node] = cur_cost + dis
                    heapq.heappush(q, (node, cur_cost+dis))
        return cost
    
    INF = float("inf")
    cost_s = [INF]*(n+1)
    cost_a = [INF]*(n+1)
    cost_b = [INF]*(n+1)
    
    cost_s = search(cost_s, s)
    cost_a = search(cost_a, a)
    cost_b = search(cost_b, b)
    
    for k in range(1, n+1):
        answer = min(cost_s[k]+cost_a[k]+cost_b[k], answer)
    
    if answer > cost_s[a]+cost_s[b]:
        answer = cost_s[a]+cost_s[b]
    
    return answer