[프로그래머스] [Python] (Level 3) 부대복귀

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

문제

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

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

 

입출력

n	roads	sources	destination	result
3	[[1,2],[2,3]]	[2, 3]	1	[1, 2]
5	[[1,2],[1,4],[2,4],[2,5],[4,5]]	[1, 3, 5]	5	[2, -1, 0]

 

제한사항

• 3 ≤ n ≤ 100,000
  • 각 지역은 정수 1부터 n까지의 번호로 구분됩니다.
• 2 ≤ roads의 길이 ≤ 500,000
  • roads의 원소의 길이 = 2
  • roads의 원소는 [a, b] 형태로 두 지역 a, b가 서로 왕복할 수 있음을 의미합니다.(1 ≤ a, b ≤ n, a ≠ b)
  • 동일한 정보가 중복해서 주어지지 않습니다.
    • 동일한 [a, b]가 중복해서 주어지지 않습니다.
    • [a, b]가 있다면 [b, a]는 주어지지 않습니다.
• 1 ≤ sources의 길이 ≤ 500
  • 1 ≤ sources[i] ≤ n
• 1 ≤ destination ≤ n

풀이방법

결론부터 말하자면, 다익스트라와 플로이드-워셜 알고리즘은 시간초과가 발생해서 BFS로 풀었습니다.

1. 그래프를 입력합니다.

2. 각각의 노드의 거리를 -1로 초기화합니다.

3. 도착지에서부터 출발지까지의 거리를 BFS를 이용해 계산합니다.

-> BFS를 사용할 수 있는 이유는 간선의 거리가 모두 1이므로 최소경로 알고리즘을 사용하지 않아도 됩니다.

4. sources에서 destination 거리를 costs배열에서 구해서 추가합니다.

 

코드

from collections import deque

def bfs(destination, graph, costs):
	# 도착지에서부터 출발지까지의 거리를 계산
    q = deque([destination])
    costs[destination] = 0
    while q:
        x = q.popleft()
        for node in graph[x]:
            if costs[node] == -1:
                costs[node] = costs[x] + 1
                q.append(node)
    return costs
    
def solution(n, roads, sources, destination):
	# 그래프 추가
    graph = [[] for _ in range(n + 1)]
    for start, end in roads:
        graph[start].append(end)
        graph[end].append(start)
    costs = [-1] * (n + 1)
    costs = bfs(destination, graph, costs)

    return [costs[s] for s in sources]

 

결론

최소거리라는 식만보고 다익스트라를 이용해 풀어서 시간초과가 발생했습니다.

이후 플로이드 워셜을 풀었는데 이것은 시간복잡도 n^3이었지만 혹시나 풀어봤습니다. 하지만 시간초과가 발생했습니다.

BFS를 이용해 시간초과가 발생하지 않았습니다.

다들 BFS로 풉시다!