강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 1로 동일합니다. 임무를 수행한 각 부대원은 지도 정보를 이용하여 최단시간에 부대로 복귀하고자 합니다. 다만 적군의 방해로 인해, 임무의 시작 때와 다르게 되돌아오는 경로가 없어져 복귀가 불가능한 부대원도 있을 수 있습니다.
강철부대가 위치한 지역을 포함한 총지역의 수 n, 두 지역을 왕복할 수 있는 길 정보를 담은 2차원 정수 배열 roads, 각 부대원이 위치한 서로 다른 지역들을 나타내는 정수 배열 sources, 강철부대의 지역 destination이 주어졌을 때, 주어진 sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열을 return하는 solution 함수를 완성해주세요. 복귀가 불가능한 경우 해당 부대원의 최단시간은 -1입니다.
from heapq import heappush, heappop
from collections import defaultdict
INF = int(1e9)
def solution(n, roads, sources, destination):
distance = [INF] * (n+1)
graph = defaultdict(list)
#그래프 만들기
for src, dest in roads:
graph[src].append(dest)
graph[dest].append(src)
#다익스트라 함수
def dijkstra(start):
q = []
distance[start] = 0
heappush(q, (0, start))
while q:
dist, now = heappop(q)
#만약 처리된 노드라면 스킵한다
if dist > distance[now]:
continue
#처리되지 않은 노드라면 다음 과정 수행
for v in graph[now]:
cost = dist + 1
if cost < distance[v]:
heappush(q, (cost, v))
distance[v] = cost
#다익스트라 수행
dijkstra(destination)
answer = []
for s in sources:
if distance[s] == INF:
answer.append(-1)
else:
answer.append(distance[s])
return answer