문제 설명
강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 1로 동일합니다. 임무를 수행한 각 부대원은 지도 정보를 이용하여 최단시간에 부대로 복귀하고자 합니다. 다만 적군의 방해로 인해, 임무의 시작 때와 다르게 되돌아오는 경로가 없어져 복귀가 불가능한 부대원도 있을 수 있습니다.
강철부대가 위치한 지역을 포함한 총지역의 수 n, 두 지역을 왕복할 수 있는 길 정보를 담은 2차원 정수 배열 roads, 각 부대원이 위치한 서로 다른 지역들을 나타내는 정수 배열 sources, 강철부대의 지역 destination이 주어졌을 때, 주어진 sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열을 return하는 solution 함수를 완성해주세요. 복귀가 불가능한 경우 해당 부대원의 최단시간은 -1입니다.
제한 사항
- 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
입출력 예
| 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] |
입출력 예 설명
입출력 예 #1
- 지역 2는 지역 1과 길로 연결되어 있기 때문에, 지역 2에서 지역 1의 최단거리는 1입니다.
- 지역 3에서 지역 1로 이동할 수 있는 최단경로는 지역 3 → 지역 2 → 지역 1 순으로 이동하는 것이기 때문에, 지역 3에서 지역 1의 최단거리는 2입니다.
- 따라서 [1, 2]를 return합니다.
입출력 예 #2
- 지역 1에서 지역 5의 최단경로는 지역 1 → 지역 2 → 지역 5 또는 지역 1 → 지역 4 → 지역 5 순으로 이동하는 것이기 때문에, 최단거리는 2입니다.
- 지역 3에서 지역 5로 가는 경로가 없기 때문에, 지역 3에서 지역 5로 가는 최단거리는 -1입니다.
- 지역 5에서 지역 5는 이동할 필요가 없기 때문에, 최단거리는 0입니다.
따라서 [2, -1, 0]을 return합니다.
풀이
최단 거리를 구하라니깐 다익스트라를 쓰면 되겠다.
효율성을 위해 dist 배열을 heap으로 만들고 이를 distHeap이라고 칭했다.
✔️ 다만 복귀 부대를 dest로 해서 모든 src에 대해 dijkstra를 적용하면 시간 초과가 뜬다.
dijkstra가 한 점에서 모든 다른 점까지 가는 최소 거리를 구하는 알고리즘이므로, 역으로 복귀 부대에서 시작해서 src까지의 거리를 구하자. 이 문제에서 그래프는 양 방향이므로 가능하다.
distHeap은 거리가 갱신된 정점이 들어가는 Heap이다.distList는 시작점에서distList[i]까지의 거리를 담은 리스트이다.
distHeap에는 시작점을 넣고,distList[시작점] = 0으로 초기화 해준다.distHeap이 empty일때까지 while문을 돈다.distHeap에서 거리가 최소인 정점을 꺼내고, 방문했다고 표시해 준다.- 해당 정점의 이웃을 돌면서, 아직 방문하지 않았으며 거리가 갱신 가능할 경우 거리를 갱신해 주고,
distHeap에 넣는다.
코드
from heapq import heappush, heappop
def makeGraph(n, roads) :
graph = dict()
for i in range(1, n+1) :
graph[i] = []
for v1, v2 in roads :
graph[v1].append(v2)
graph[v2].append(v1)
return graph
def dijkstra(src, n, graph) :
distHeap = [(0, src)]
distList = [float("inf")] * (n+1)
visited = [False] * (n+1)
distList[src] = 0
while distHeap :
curDist, cur = heappop(distHeap)
visited[cur] = True
for adjacent in graph[cur] :
if not visited[adjacent] and distList[adjacent] > curDist + 1 :
distList[adjacent] = curDist + 1
heappush(distHeap, (curDist+1, adjacent))
return distList
def solution(n, roads, sources, destination):
answer = []
graph = makeGraph(n, roads)
distList = dijkstra(destination, n, graph)
for source in sources :
answer.append(-1 if distList[source]==float("inf") else distList[source])
return answer
반갑습니다.