프로그래머스 - 부대복귀 (Python)

Lv3, bfs

문제 설명

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

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

입출력 예시

풀이

부대원의 각 위치부터 destination 까지 for을 돌면서 bfs를 수행하면 시간초과가 발생했다.

-> destination으로 부터 각 지점까지 하면?

from collections import deque

def solution(n, roads, sources, end):
    graph = [[] for _ in range(n + 1)]

    for road in roads:
        graph[road[0]].append(road[1])
        graph[road[1]].append(road[0])

    dist = [-1] * (n + 1)

    # 생각을 바꿔보자
    # destination -> start로 가도 된다.

    def bfs():
        q = deque()
        dist[end] = 0
        q.append(end)

        while q:
            now = q.popleft()

            for v in graph[now]:
                if dist[v] == -1:
                    dist[v] = dist[now] + 1
                    q.append(v)

    bfs()
    answer = []
    for source in sources:
        answer.append(dist[source])

    return answer

dist값을 지정하는 부분에서 min()을 통해 최소를 넣어줘야하는게 아닐까 하는 의구심이 드는데 그냥 문제가 정답처리를 한다.

Lv3 이라기엔 난이도가 매우 쉬운 편