[Python/파이썬] 백준 1238번 파티

1238번: 파티

 

문제

N개의 숫자로 구분된 각각의 마을에 한 명의 학생이 살고 있다.

어느 날 이 N명의 학생이 X (1 ≤ X ≤ N)번 마을에 모여서 파티를 벌이기로 했다. 이 마을 사이에는 총 M개의 단방향 도로들이 있고 i번째 길을 지나는데 Ti(1 ≤ Ti ≤ 100)의 시간을 소비한다.

각각의 학생들은 파티에 참석하기 위해 걸어가서 다시 그들의 마을로 돌아와야 한다. 하지만 이 학생들은 워낙 게을러서 최단 시간에 오고 가기를 원한다.

이 도로들은 단방향이기 때문에 아마 그들이 오고 가는 길이 다를지도 모른다. N명의 학생들 중 오고 가는데 가장 많은 시간을 소비하는 학생은 누구일지 구하여라.

입력

첫째 줄에 N(1 ≤ N ≤ 1,000), M(1 ≤ M ≤ 10,000), X가 공백으로 구분되어 입력된다. 두 번째 줄부터 M+1번째 줄까지 i번째 도로의 시작점, 끝점, 그리고 이 도로를 지나는데 필요한 소요시간 Ti가 들어온다. 시작점과 끝점이 같은 도로는 없으며, 시작점과 한 도시 A에서 다른 도시 B로 가는 도로의 개수는 최대 1개이다.

모든 학생들은 집에서 X에 갈수 있고, X에서 집으로 돌아올 수 있는 데이터만 입력으로 주어진다.

출력

첫 번째 줄에 N명의 학생들 중 오고 가는데 가장 오래 걸리는 학생의 소요시간을 출력한다.

코드

• 플로이드 워셜 풀이

import sys
input = sys.stdin.readline

INF = int(1e9)

n, m, x = map(int, input().split())

road = [[INF for _ in range(n + 1)] for _ in range(n + 1)]
for i in range(n + 1):
    road[i][i] = 0  # 자기 자신으로 가는 길은 0으로 초기화
    
for _ in range(m):
    s, e, t = map(int, input().split())
    road[s][e] = t
    
# 플로이드 워셜  
for k in range(1, n + 1):
    for i in range(1, n + 1):
        for j in range(1, n + 1):
            road[i][j] = min(road[i][j], road[i][k] + road[k][j])
            
answer = 0
for c in range(1, n + 1):
    answer = max(answer, road[x][c] + road[c][x])

print(answer)

N의 입력이 1000까지 들어올 수 있으므로 시간복잡도가 O(N^3)인 플로이드 워셜을 사용하면 시간초과가 발생한다. 따라서 아래와 같이 다익스트라로 다시 풀었다. 

 

• 다익스트라 풀이

import sys, heapq
input = sys.stdin.readline

INF = int(1e9)

n, m, x = map(int, input().split())

road = [[] for _ in range(n + 1)]
for _ in range(m):
    s, e, t = map(int, input().split())
    road[s].append((e, t))

def dijkstra(start):
    q = []
    distance = [INF] * (n + 1)
    # 자기자신으로 가는 최단 경로는 0으로 설정하여 큐에 삽입
    heapq.heappush(q, (0, start))
    distance[start] = 0
    
    while q:
        dist, now = heapq.heappop(q)
        
        if distance[now] < dist:    # 이미 처리된 노드
            continue
        for e, t in road[now]:
            cost = dist + t
            # 현재 노드를 거치는게 더 짧은 경로일 경우
            if cost < distance[e]:
                distance[e] = cost
                heapq.heappush(q, (cost, e))
                
    return distance

# 각 마을의 학생들이 x번 마을을 오고 가는 값 중 최댓값 구하기
answer = 0
x_to_each = dijkstra(x)

for i in range(1, n + 1):
    if i == x:
        continue
    answer = max(answer, x_to_each[i] + dijkstra(i)[x])
    
print(answer)

다익스트라 알고리즘을 수행하는 함수를 작성하여 한 노드에서 다른 노드로 가는 최단 거리를 담은 리스트를 리턴하도록 하였다. 이 함수를 이용하여 우선 X번 마을에서 다른 모든 마을로 가는 최단 거리를 x_to_each라는 리스트에 담아두고 모든 마을에 대해 다익스트라 함수를 적용하여 각 마을에서 X번 마을로 가는 최단 거리를 구한다. 이 최단 거리와 x_to_each에 저장된 X번 마을에서 각 마을로 가는 최단 거리를 구한 값 중 최댓값을 구하면 문제의 답이다.