코딩 테스트(Python) 부대 복귀

문제 설명

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

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

 

from collections import deque

def solution(n, roads, sources, destination):
    answer = []  # 결과를 저장할 리스트
    # 각 노드에 연결된 노드를 저장할 그래프 초기화
    graph = [[] for i in range(n+1)]
    
    # 도로 정보(roads)를 바탕으로 그래프를 구성
    for i, j in roads:
        graph[i].append(j)  # i에서 j로 가는 도로 추가
        graph[j].append(i)  # j에서 i로 가는 도로 추가 (양방향)
    
    # 각 노드에 대한 최단 거리 초기화, -1은 아직 방문하지 않은 상태를 의미
    distances = [-1 for i in range(n+1)]
    distances[destination] = 0  # 목적지의 거리는 0으로 설정
    
    # BFS를 위한 큐 초기화
    q = deque()
    q.append([destination, 0])  # 큐에 목적지와 거리(0)를 삽입
    
    # BFS 탐색
    while q:
        current, dist = q.popleft()  # 큐에서 현재 노드와 그에 해당하는 거리를 꺼냄
        for i in graph[current]:  # 현재 노드와 연결된 노드들을 탐색
            if distances[i] == -1:  # 아직 방문하지 않은 노드라면
                distances[i] = dist + 1  # 거리를 1 증가시켜 저장
                q.append([i, dist + 1])  # 큐에 노드와 새로운 거리를 삽입
    
    # 각 출발지에서 목적지까지의 최단 거리를 답 리스트에 저장
    answer = [distances[i] for i in sources]
    
    return answer  # 출발지들에 대한 최단 거리 리스트 반환