[programmers/py] 부대복귀

부대복귀

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

문제 설명

강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 1로 동일합니다. 임무를 수행한 각 부대원은 지도 정보를 이용하여 최단시간에 부대로 복귀하고자 합니다. 다만 적군의 방해로 인해, 임무의 시작 때와 다르게 되돌아오는 경로가 없어져 복귀가 불가능한 부대원도 있을 수 있습니다.

강철부대가 위치한 지역을 포함한 총지역의 수 n, 두 지역을 왕복할 수 있는 길 정보를 담은 2차원 정수 배열 roads, 각 부대원이 위치한 서로 다른 지역들을 나타내는 정수 배열 sources, 강철부대의 지역 destination이 주어졌을 때, 주어진 sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열을 return하는 solution 함수를 완성해주세요. 복귀가 불가능한 경우 해당 부대원의 최단시간은 -1입니다.

풀이

각 노드에서 갈 수 있는 노드를 보관 후 목적지 -> 출발지로 이동하며 각 출발지까지 거리를 저장합니다.

저장된 거리에서 출발지->목적지까지 거리를 answer에 추가합니다.

from collections import deque
def solution(n, roads, sources, destination):
    answer = []
    visited = [False for _ in range(n+1)]
    graph = [[]*(n+1) for _ in range(n+1)]
    
    for x, y in roads:
        graph[x].append(y)
        graph[y].append(x)
    
    dis = [-1] * (n+1)
    dq = deque()
    dq.append([destination, 0])
    
    while dq:
        now, cnt = dq.popleft()
        
        if visited[now]:
            continue 
        
        dis[now] = cnt
        visited[now] = True
        
        for n in graph[now]:
            dq.append([n, cnt+1])
                
    for s in sources:
        answer.append(dis[s])
    return answer