[Python] 백준 1238번 - 파티 (골드 3)

혼자 힘으로 풀었는가? O

알고리즘 분류
 - 최단 경로
 - 다익스트라

 

https://www.acmicpc.net/problem/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명의 학생들 중 오고 가는데 가장 오래 걸리는 학생의 소요시간을 출력한다.

 

---

 

다익스트라 알고리즘을 이용하여 최단경로를 N번 구하는 문제이다.

 

나는 구분하기 편하게 1~N까지의 가는 리스트인 go 리스트와

 

X에서 1~N으로 돌아가는 back 리스트로 구분하여 문제를 풀었다.

 

import sys
import heapq

# 표준 입력에서 읽기 위해 입력 방식을 변경
input = sys.stdin.readline

# 노드 수, 간선 수, 목적지 노드 X를 읽음
n, m, x = map(int, input().split())

# 무한대로 그래프를 초기화하고 인접 리스트를 설정
INF = int(1e9)
data = [[] for _ in range(n+1)]
for i in range(m):
    a, b, w = map(int, input().split())  # 각 간선과 그 가중치를 읽음
    data[a].append([b, w])  # 인접 리스트에 간선 추가
    
# 최단 경로를 찾기 위한 다익스트라 알고리즘 구현
def dijkstra(start):
    q = []  # 최소 힙을 위한 우선순위 큐
    tmp = [INF for _ in range(n+1)]  # 거리를 INF로 초기화
    
    heapq.heappush(q, [0, start])  # 시작 노드를 거리 0으로 푸시
    tmp[start] = 0
    while q:
        weight, now = heapq.heappop(q)  # 가장 작은 거리를 가진 노드를 팝
        for next, w2 in data[now]:  # 모든 인접 노드를 탐색
            if tmp[now] + w2 < tmp[next]:  # 더 짧은 경로를 찾으면 거리 업데이트
                tmp[next] = tmp[now] + w2
                heapq.heappush(q, [w2, next])  # 우선순위 큐에 업데이트된 거리를 푸시
    return tmp

# 왕복 거리의 최대값을 저장할 result 변수 초기화
result = 0

# X에서 모든 다른 노드로의 최단 경로를 계산
back = dijkstra(x)

# 모든 노드에서 X로의 최단 경로를 저장할 리스트 초기화
go = [[] for _ in range(n+1)]
for i in range(1, n+1):
    if i == x:  # 현재 노드가 X이면 건너뜀
        continue
    go[i] = dijkstra(i)  # 노드 i에서 X로의 최단 경로를 찾음
    
# X로 가고 돌아오는 왕복 거리의 최대값을 계산
for i in range(1, n+1):
    if i == x:  # 현재 노드가 X이면 건너뜀
        continue
    result = max(result, go[i][x] + back[i])  # 왕복 거리의 최대값으로 결과 업데이트
    
# 왕복 거리의 최대값을 출력
print(result)

 

풀고보니 해당 코드는 상당히 비효율적인 정답인 것 같다.

 

힙큐에 넣을 때 사용하지도 않는 weight를 넣고 있고

 

최대값을 찾기 위해서도 불필요한 반복문을 사용하는 등 개선의 여지가 많은 부분이다.

 

import sys
import heapq

input = sys.stdin.readline

n, m, x = map(int, input().split())
graph = [[] for _ in range(n+1)]
for _ in range(m):
    a, b, c = map(int, input().split())
    graph[a].append((b, c))

def dijkstra(start):
    distance = [1e9] * (n + 1)
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0
    while q:
        dist, now = heapq.heappop(q)
        if distance[now] < dist:
            continue
        for i in graph[now]:
            cost = dist + i[1]
            if cost < distance[i[0]]:
                distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))
    return distance

result = 0
# X에서 모든 마을로 가는 최단 거리 계산
distance_from_x = dijkstra(x)
# 모든 마을에서 X로 가는 최단 거리 계산
for i in range(1, n+1):
    distance_to_x = dijkstra(i)
    # 왕복 거리 합이 최대가 되는 값을 찾음
    result = max(result, distance_from_x[i] + distance_to_x[x])

print(result)

 

아래는 내 코드이며 위는 개선된 코드이다.

 

속도가 배 이상으로 빨라진 것을 알 수 있다.

 

너무 어렵다...