[SWEA-파이썬] 1795: 인수의 생일 파티

문제 보러 가기!

N개의 집(정점)이 주어지고, 집과 집을 연결하는 도로(간선)가 주어진다. 도로는 단방향 간선이고, 이동하는 데 걸리는 시간이 주어진다.
도로는 모든 집들 간의 이동이 가능하게 했다.
모든 집에서 생일을 맞은 인수의 집으로 오고 가는데, 이때 오가는 시간이 가장 오래 걸리는 집이 얼마나 걸리는지 구하는 문제이다.

그러면 인수네 집 -> 다른 모든 집으로 가는 데 드는 최소 시간도, 다른 모든 집 각각에서 인수네 집으로 가는 데 드는 최소 시간도 모두 알아야 한다.

인수네 집에서 다른 모든 집들 가는 최단 경로는 다익스트라로 할 수 있다는 건 처음부터 알았다.
다른 각각의 집에서 인수네 집으로 가는 최단 경로는 어떻게 구할까.. 그게 고민이었다.

처음에는 마음 편하게 ! 모든 정점 간의 최단 경로를 찾아주는 플로이드-워샬 알고리즘으로 문제를 풀었다.
답은.. 나오는데, 시간 초과가 나왔다.. O(N^3)... 너무 시간이 오래 걸린다.ㅠㅠ

다른 각각의 집에서 인수네 집으로 가는 최단 경로를 어떻게 구할지 좀 더 생각했다.
도로의 방향을 반대로 받은 후에, 인수네 집에서 다른 집으로 가는 최단 경로를 구하면 되는 것이었다 !!! 🎉🎉 

간선 정보를 받을 때 원래 방향대로 받은 리스트(origin_abj)를 만들고, 반대 방향으로 받은 리스트(opposite_abj)를 만들었다.
이 두 간선 정보를 가지고 인수네 집에서 출발하는 다익스트라를 각각 돌려서, 인수네에서 다른 집에 가는 최단 경로, 다른 집에서 인수네로 오는 최단 경로를 구했다. 

이 둘을 더한 합이 가장 큰 경우를 답으로 했다.

그랬더니 성공!! (다익스트라 알고리즘의 시간 복잡도는 O(ElogE))

import heapq

def dijikstra(start, abj):
    # 비용 담긴 리스트
    cost = [987654321] * (V + 1)
    # 방문 체크
    visited = [False] * (V + 1)

    # 출발점 초기화
    Q = []
    heapq.heappush(Q, (0, start))  # 출발점에서 어떤 집까지 가는 데 드는 비용, 그 집 번호

    while Q:
        curr_cost, curr_house = heapq.heappop(Q)
        if visited[curr_house]:
            continue
        visited[curr_house] = True
        cost[curr_house] = curr_cost

        # 그 집 근처의 비용 업데이트
        for i in range(1, V + 1):
            if visited[i] == False:  # 방문 안 했으면
                cost[i] = min(cost[i], cost[curr_house] + abj[curr_house][i])
                heapq.heappush(Q, (cost[i], i))
    return cost


T = int(input())

for tc in range(1, T + 1):
    V, E, insu_house = map(int, input().split())
    origin_abj = [[987654312] * (V + 1) for _ in range(V + 1)]
    opposite_abj = [[987654321] * (V + 1) for _ in range(V + 1)]

    for _ in range(E):
        s, e, cost = map(int, input().split())
        origin_abj[s][e] = cost
        opposite_abj[e][s] = cost

    # 자기 위치 0으로
    for i in range(V+1):
        origin_abj[i][i] = 0
        opposite_abj[i][i] = 0

    origin_res = dijikstra(insu_house, origin_abj)
    oppo_res = dijikstra(insu_house, opposite_abj)

    # 계산
    ans  = 0
    for i in range(1,V+1):
        ans = max(ans, origin_res[i] + oppo_res[i])

    print(f'#{tc} {ans}')