[Python 문제 풀이] 프로그래머스 '부대 복귀'

거리를 기록해놓자

 

접근

 node와 link의 수가 꽤 많다. 그래서 그런가 그냥 BFS 적용하면 시간초과가 발생한다. 나처럼 쉬운문제부터 올라온 사람은 시간초과 문제를 겪을 것이다. 이 경우 기존에 방문했던 node들의 거리를 기록해놓고 다른 경로를 탐색할 때 활용해야 한다.

 이러한 접근 방법은 동적계획법과 같다. A->B->C의 경로는 A->B와 B->C의 경로라는 부분문제로 나눌 수 있고 B까지의 경로를 사전에 계산해놓았다면 B->C만 계산하면 된다.

 

의사코드

주어진 배열을 인접리스트로 바꾼다.

노드의 수를 길이로 하는 costs배열을 만들고 모든 원소를 -1로 초기화한다.

출발점이 될 costs[destination]에 0을 대입한다.

destination에서 출발해 인접한 노드의 거리를 +1씩 추가하여 costs에 기록한다.

인접한 노드들을 대기열에 추가한다.

대기열이 빌 때까지 대기열에 추가된 노드들에 대해서 반복한다.

sources에 대해서 반복하며 costs에 저장된 최단거리를 정답에 기록한다.

 

구현

from collections import deque
def solution(n, roads, sources, destination):
    answer = []
    graph = [[] for _ in range(n + 1)]
    costs = [-1 for _ in range(n + 1)]
    costs[destination] = 0
    queue = deque([destination])
    for n1, n2 in roads :
        graph[n1].append(n2)
        graph[n2].append(n1)
    while queue :
        x = queue.popleft()
        for node in graph[x] :
            if costs[node] == -1 :
                queue.append(node) 
                costs[node] = costs[x] + 1
    for s in sources :
        answer.append(costs[s])
    return answer