강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 1로 동일합니다. 임무를 수행한 각 부대원은 지도 정보를 이용하여 최단시간에 부대로 복귀하고자 합니다. 다만 적군의 방해로 인해, 임무의 시작 때와 다르게 되돌아오는 경로가 없어져 복귀가 불가능한 부대원도 있을 수 있습니다.
강철부대가 위치한 지역을 포함한 총지역의 수 n, 두 지역을 왕복할 수 있는 길 정보를 담은 2차원 정수 배열 roads, 각 부대원이 위치한 서로 다른 지역들을 나타내는 정수 배열 sources, 강철부대의 지역 destination이 주어졌을 때, 주어진 sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열을 return하는 solution 함수를 완성해주세요. 복귀가 불가능한 경우 해당 부대원의 최단시간은 -1입니다.
import heapq
INF = int(1e9)
def solution(n, roads, sources, destination):
answer = []
graph = [[] for i in range(n +1)]
for a,b in roads:
graph[a].append(b)
graph[b].append(a)
visited = dij(destination,graph,n)
for i in sources:
if visited[i] == INF :
answer.append(-1)
else:
answer.append(visited[i])
return answer
def dij(start, graph,n):
q = []
heapq.heappush(q,(0,start))
visited = [INF for i in range(n +1)]
visited[start] = 0
while q:
cost, start = heapq.heappop(q)
for i in graph[start]:
if visited[i] != INF:
continue
visited[i] = cost +1
heapq.heappush(q,(cost+1,i))
return visited
사실 이문제는 처음에 시도했을때 11~15번까지가 시간 복잡도 초과가 떴다.(실패)
고민을했다.
출발지와 도착지점이 같을 때 먼저 제외를 해야할까? 혹은 어떻게 sources를 안 탐색할수가 있을까?
하지만 극단적인 상황을 제외하고는 해당 방법들은 소용이 없었다.
그래서 생각해낸 풀이
만약 대전 -> 서울 , 춘천 -> 서울, 울산 -> 서울 을 간다고하면
어느 역에 물어봐야 해당 모든 시간표를 알까?
정답은 도착 지점이다.
시간복잡도가 현재 dij를 돌며 nlongn * n * s 이니
도착 지점에서 가는 최소 거리를 구한 지도를 받아와서 각 지점에 대해 S*O(1)으로 구해주는 것이다.
해당 문제로 풀었더니 시간 초과가 풀렸고 모두 정답 처리가 되었다.
역발상으로 푸는 재미있는 문제였다.