[프로그래머스] 부대복귀 파이썬

문제링크: https://school.programmers.co.kr/learn/courses/30/lessons/132266

난이도: LV 3

문제해결 아이디어

• destination을 q에 넣고 갈 수 있는 모든 최단 경로를 구한다.
• visited에는 인덱스별로 최단 경로값을 바인딩
• source에 있는 부대원들의 번호를 visited를 통해 값을 리턴한다.

from collections import deque

def solution(n, roads, sources, destination):
    visited = [-1] * (n+1)
    graph = [[] for _ in range(n+1)]
    for a,b in roads:
        graph[a].append(b)
        graph[b].append(a)
    
    q = deque()
    q.append(destination)
    visited[destination] = 0
    
    while q:
        now = q.popleft()
        
        for node in graph[now]:
            if visited[node] == -1:
                visited[node] = visited[now] + 1
                q.append(node)
    
    return [visited[i] for i in sources]