[PS] BOJ 1238 / 파티

[PS] BOJ 1238 / 파티
문제 링크: https://www.acmicpc.net/problem/1238
Thumbnail: Photo by Adi Goldstein (Unsplash)

파티를 즐기기 위해서는 최단 거리를 구해야 합니다!

풀이

다익스트라(Dijkstra)!

다익스트라 알고리즘을 활용하는 문제입니다. 한가지 유의할 점은, 유향 그래프이기 때문에 각 학생이 사는 마을에서 파티가 열리는 $X$번 마을로 가는 거리와, $X$번 마을에서 다시 집으로 돌아가는 거리는 다를 수 있습니다.

따라서, 왕복 거리는 다음과 같이 계산하면 됩니다.

to_X = dijkstra(X)  # X에서 다른 모든 정점까지의 최단 거리

max_distance = 0
for i in range(1, N + 1):
    distance = to_X[i] + dijkstra(i)[X]  # X에서 i까지의 최단 거리 + i에서 X까지의 최단 거리
    max_distance = max(max_distance, distance)

결과 계산하기

다익스트라 알고리즘은 우선순위 큐 기반으로 구현했습니다. 상세한 풀이는 이전 문제를 참고해 주세요 😉

전체 코드

from heapq import heappop, heappush

input = open(0).readline
N, M, X = map(int, input().split())
INF = int(1e9)

# 간선 입력
edges = [[] for _ in range(N + 1)]
for _ in range(M):
    u, v, cost = map(int, input().split())
    edges[u].append((v, cost)) # 유향 그래프
edges

# Dijkstra
def dijkstra(start):
    distance = [INF for _ in range(N + 1)]
    distance[start] = 0
    pq = []
    heappush(pq, (0, start))  # (비용, 정점)
    
    while pq:
        cost, u = heappop(pq)
        if distance[u] < cost:  # 이미 더 낮은 비용으로 갈 수 있다면 굳이 반영하지 않는다.
            continue
        for v, c in edges[u]:
            new_dist = cost + c
            if new_dist < distance[v]:
                distance[v] = new_dist
                heappush(pq, (distance[v], v))
    
    return distance

to_X = dijkstra(X)  # X에서 다른 모든 정점까지의 최단 거리

max_distance = 0
for i in range(1, N + 1):
    distance = to_X[i] + dijkstra(i)[X]  # i에서 X까지의 최단 거리
    max_distance = max(max_distance, distance)

# 결과 출력
print(max_distance)  # X에서 다른 모든 정점까지의 최단 거리 중 최대값

solution.py