[백준 BOJ_1238] 파티 Python 풀이

출처: 백준 온라인 저지

문제

풀이

다익스트라(dijkstra)를 이용하여 구현해주었습니다.

간선들을 확인하여 인접리스트를 만들어주었습니다. 단방향이기 때문에 서로에게 추가해주는 것이 아닌 출발점을 key로 가지고, 도로의 길이와 도착점의 tuple을 heappush를 이용하여 list에 넣어주었습니다. 여기서 도로가 짧은 순으로 넣어주어야하기 때문에 꼭 도로의 길이, 도착점의 순으로 tuple을 만들어 넣어주어야 합니다.

그 다음, dijkstra 함수를 만들어주었습니다. 시작점인 src와 인접리스트 adj를 인자로 받고, dist를 float(“inf”)로 초기화해줍니다. dist에는 출발점인 src에서부터 나머지 노드들까지의 최단거리를 담아주는 list입니다. 본격적으로 while문을 돌아주기 전에 자기자신과의 거리는 0이므로 dist의 출발점 src의 거리를 0으로 저장해줍니다. hq에는 누적해줄 도로의 길이(0), 출발점(src)를 넣어주고 while문을 시작합니다.

heappop을 해주게 되면 hq안에서 누적된 도로가 짧은 순으로 나오게 됩니다. heappop해준 도로 길이(w)와 노드(v) 중 해당 노드의 인접리스트를 돌아줍니다. 새로운 도로 길이(nw)에 여태까지 누적해준 도로 길이(w)를 더해주고, 그 길이가 현재 새로운 노드의 최단거리 값(dist[nv])보다 짧다면, 갱신해주고 hq에 위와 같은 (도로 길이, 노드)의 tuple을 heappush해줍니다. 그렇게 hq를 다 돌아주고 난 뒤 dist를 return해줍니다.

일단 dijkstra 함수는 출발점에서 모든 노드까지의 최단거리를 담는 list를 return하기 때문에 파티를 하는 마을(X)에서 각자의 마을로 돌아가는 경우로 먼저 dist를 초기화해줍니다. 그리고 마을을 가르키는 숫자가 1부터 시작하므로 편의상 N+1 크기의 dist를 만들어주었기 때문에 index 0에는 inf가 들어가 있으므로 0으로 초기화해줍니다. 나중에 max함수를 쓸 때 index 0에 있는 inf값이 걸리지 않게끔하기 위함입니다. 그 다음 1부터 N까지 출발점으로 dijkstra함수를 돌려주며 to_x에 저장해줍니다. 파티를 하는 마을(X)까지의 최단거리(to_x[X])를 dist[i]에 더해줍니다. 단, i와 X가 같을 경우는 자기 자신의 마을에서 파티를 하는 경우이기 때문에 따로 더해주지 않습니다.

마지막에는 dist의 최대값을 max함수를 사용하여 출력해줍니다.

코드

import sys
from collections import defaultdict
from heapq import heappush, heappop
input = sys.stdin.readline
N, M, X = map(int,input().split())
adj = defaultdict(list)
for _ in range(M):
  v1, v2, w = map(int,input().split())
  heappush(adj[v1],(w, v2))
def dijkstra(src, adj):
  dist = [float("inf") for _ in range(N+1)]
  dist[src] = 0
  hq = [(0, src)]
  while hq:
    w, v = heappop(hq)
    for nw, nv in adj[v]:
      nw += w
      if dist[nv] > nw:
        dist[nv] = nw
        heappush(hq, (nw, nv))
  return dist
dist = dijkstra(X, adj)
dist[0] = 0
for i in range(1, N+1):
  if i == X:
    dist[i] = 0
    continue
  to_x = dijkstra(i, adj)
  dist[i] += to_x[X]
print(max(dist))