[최단거리] 부대 복귀 (파이썬)

문제 설명

https://school.programmers.co.kr/learn/courses/30/lessons/132266

프로그래머스

문제 풀이

• 주어진 양방향 그래프 경로들을 통해 source에 담겨진 원소들이 destination에 도착하는 최단 경로를 각각 구하는 문제
• 양방향 그래프이므로 source에서 출발한 점들이 des에 도착하는 거리나 des에서 각 점들에 도착하는 거리나 똑같음.
• start를 des로 두고 경로들을 탐색하면 됨

 

소스 코드

import heapq

def solution(n, roads, sources, destination):
    gr = [[] for i in range(n+1)]
    INF = int(1e9)
    distance = [INF]*(n+1)
    distance[destination]=0
    answer = []
    for a,b in roads:
        gr[a].append(b)
        gr[b].append(a)
    
    q = []
    heapq.heappush(q,(0,destination))
    while(q):
        d, cur = heapq.heappop(q)
        for next in gr[cur]:
            if distance[next]<=d+1:
                continue
            distance[next]=d+1
            heapq.heappush(q,(d+1,next))
    for s in sources:
        if distance[s]!=INF:
            answer.append(distance[s])
        else:
            answer.append(-1)
    return answer
    

print(solution(5,   [[1, 2], [1, 4], [2, 4], [2, 5], [4, 5]],   [1, 3,5], 5))
        # 총 지역 수 , 왕복 가능한 길 정보,   각 부대원 위치 지역, 도착 지점
        # source의 순서대로 강철부대로 복귀 할 수 있는 최단 시간