백준 1238번 파티 파이썬

리뷰

처음엔 파티 장소까지 갔던 인덱스를 구해 다시 돌아오면 되는 줄 알았지만 문제를 자세히 읽어보니 왕복 거리가 가장 큰 수를 출력해 내는 것이였다. 문제를 잘 읽자...

 

문제 풀이

1. 우선 모든 마을에서 x번 마을까지의 거리를 모두 구해준 뒤 거리 정보의 리스트를 리턴해준다. 난 go에 저장했다.
2. 이제 for문을 통해 x에서 모든 마을로 가는 경우의 수를 구해준 뒤 back이라는 리스트에 저장해준다.
3. go와 back을 더한 값이 가장 높은 경우가 왕복거리가 제일 큰 학생의 거리가 된다.

 

참고 사항

1. 함수를 사용하지 않으면 코드가 너무 길어지기에 적절히 bool 인자를 섞어 함수로 구현하자
2. dist를 INF로 초기화 해줄때 0번 인덱스는 0으로 바꿔줘야 max를 사용하기 편하다.
3. go와 back을 비교하기 쉽게 back에도 0번 인덱스에 0을 넣어주면 좋다.

 

정답 코드

import sys
import heapq

def d(start, back, index):
    dist = [INF] * (n + 1)
    dist[0] = 0
    heap = [[0, start]]
    dist[start] = 0
    while heap:
        cv, cn = heapq.heappop(heap)
        if dist[cn] != cv: continue
        for nv, nn in e[cn]:
            if dist[nn] > cv + nv:
                dist[nn] = cv + nv
                heapq.heappush(heap, [dist[nn], nn])
    if back:
        return dist[index]
    else:
        return dist

INF = sys.maxsize
n, m, x = map(int, sys.stdin.readline().split())
e = [[] for _ in range(n + 1)]
for _ in range(m):
    a, b, c = map(int, sys.stdin.readline().split())
    e[a].append([c, b])
go = d(x, 0, 0)
back = [0] * (n + 1)
for i in range(1, n + 1):
    back[i] = d(i, 1, x)
max_dist = -1
for i in range(1, n + 1):
    max_dist = max(max_dist, go[i] + back[i])
print(max_dist)