[프로그래머스] 합승 택시 요금(Python)

• 2021.02.14 수정
  
  

1. 요약

목적지가 다른 2명이 같은 위치에서 택시를 탄다. 그리고 특정 지점에서 하차한 이후, 각자 서로 다른 택시를 타고 본인의 목적지에 도착한다. 이 때, 나올 수 있는 총 요금들 중 가장 작은 값을 구해야한다.

2. 아이디어

가중치가 있는 그래프에서 최단 거리(최소 요금)를 구하기 위해, 다익스트라 사용

A와 B가 헤어지는 지점을 K라고 하면,
요금 = dijkstra(출발점, K) + dijkstra(K, A의 도착지점) + dijkstra(K, B의 도착지점)

주어진 모든 지점(e.g. 출발점, A의 도착지점, B의 도착지점, 그 외 나머지 지점들)을 하나씩 K에 대입하여 나온 요금들을 비교한다.

3. 코드

2021.02.14 수정(최초 작성된 코드로 제출해도 정답처리 됩니다.)

수정 내용

• cost = INF로 변경,
• solution()의 for문 안에 있는 if문 삭제

from heapq import heappop, heappush

INF = int(1e9)
graph = [[]]

def preprocess(n, fares):
    global graph
    graph = [[] for i in range(n+1)]

    for fare in fares:
        src, dst, cost = fare[0], fare[1], fare[2]
        graph[src].append([dst, cost])
        graph[dst].append([src, cost])

def dijkstra(src, dst):
    global graph
    n = len(graph)
    table = [INF for i in range(n)]
    table[src] = 0
    pq = [[0, src]]

    while pq:
        w, x = heappop(pq)

        if table[x] < w: continue

        for item in graph[x]:
            nx, ncost = item[0], item[1]
            ncost += w
            if ncost < table[nx]:
                table[nx] = ncost
                heappush(pq, [ncost, nx])
    
    return table[dst]

def solution(n, s, a, b, fares):
    preprocess(n, fares)
    cost = INF

    for i in range(1, n+1):
        cost = min(cost, dijkstra(s, i) + dijkstra(i, a) + dijkstra(i, b))
    
    return cost

---

최초 작성된 코드

from heapq import heappop, heappush

INF = int(1e9)
graph = [[]]

def preprocess(n, fares):
    global graph
    graph = [[] for i in range(n+1)]

    for fare in fares:
        src, dst, cost = fare[0], fare[1], fare[2]
        graph[src].append([dst, cost])
        graph[dst].append([src, cost])

def dijkstra(src, dst):
    global graph
    n = len(graph)
    table = [INF for i in range(n)]
    table[src] = 0
    pq = [[0, src]]

    while pq:
        w, x = heappop(pq)

        if table[x] < w: continue

        for item in graph[x]:
            nx, ncost = item[0], item[1]
            ncost += w
            if ncost < table[nx]:
                table[nx] = ncost
                heappush(pq, [ncost, nx])
    
    return table[dst]

def solution(n, s, a, b, fares):
    preprocess(n, fares)
    cost = dijkstra(s, a) + dijkstra(s, b)

    for i in range(1, n+1):
        if s != i:
            cost = min(cost, dijkstra(s, i) + dijkstra(i, a) + dijkstra(i, b))
    
    return cost