백준 18352번 특정 거리의 도시 찾기 파이썬

리뷰

다익스트라를 사용해 푸는 문제, 메모리 시간 살발하다 C++로도 풀어봐야되는데 엄두가 안나네..

 

문제 풀이

1. n, m, k, x값을 모두 받아온다.
2. 각노드의 간선을 나타낼 빈 리스트를 n + 1개로 초기화 해준다.
3. 각 노드까지의 거리를 나타낼 리스트를 INF값을 갖는 n + 1개로 초기화 해준다.
4. m개의 간선 정보를 각 노드에 추가해 준다. 이때 거리는 1로 고정이므로 [1, 다음 노드]의 형식으로 추가해 줬다.
5. 힙에 시작 노드와 거리 0을 넣어주고 현재 노드에서 현재 노드까지의 거리는 0으로 설정해 준다.
6. while루프를 돌며 각 간선까지의 거리 정보를 최신화 해주고 힙에 추가해주는 작업을 해준다.
7. 루프 종료 후 거리가 k인 도시가 있다면 모두 출력해 주고 없다면 -1을 출력해 준다.

 

참고 사항

while루프 내에서 변수 명이 헷갈리기 쉽다. 조심하자

 

 

정답 코드

import sys
import heapq

INF = sys.maxsize
n, m, k, x = map(int, sys.stdin.readline().split())
edge = [[] for _ in range(n + 1)]
dist = [INF] * (n + 1)
for _ in range(m):
    a, b = map(int, sys.stdin.readline().split())
    edge[a].append([1, b])
heap = [[0, x]]
dist[x] = 0
while heap:
    cd, cn = heapq.heappop(heap)
    if dist[cn] != cd: continue
    for nd, nn in edge[cn]:
        if dist[nn] > cd + nd:
            dist[nn] = cd + nd
            heapq.heappush(heap, [dist[nn], nn])
if dist.count(k):
    for i in range(n + 1):
        if dist[i] == k:
            print(i)
else:
    print(-1)