부대 복귀
링크
[ 문제에 대한 이해 ]
강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 1로 동일합니다. 임무를 수행한 각 부대원은 지도 정보를 이용하여 최단시간에 부대로 복귀하고자 합니다. 다만 적군의 방해로 인해, 임무의 시작 때와 다르게 되돌아오는 경로가 없어져 복귀가 불가능한 부대원도 있을 수 있습니다.
강철부대가 위치한 지역을 포함한 총지역의 수n, 두 지역을 왕복할 수 있는 길 정보를 담은 2차원 정수 배열roads, 각 부대원이 위치한 서로 다른 지역들을 나타내는 정수 배열sources, 강철부대의 지역destination이 주어졌을 때, 주어진sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열을 return하는 solution 함수를 완성해주세요. 복귀가 불가능한 경우 해당 부대원의 최단시간은 -1입니다.
- 각 위치에서 강철부대로 복귀할 수 있는 최단 시간을 구한다.
- 길이 없는 경우도 존재한다. 이럴 때는 -1을 출력한다.
- 주어지는
roads는[a, b]형태로 주어지며 항상 왕복 가능하다는 것을 뜻한다.
문제의 예시를 보면 알 수 있지만, 적군의 방해로 인해 되돌아오는 경로가 없어져 라는 것에 큰 의미를 둘 필요는 없다. 단순히 이는 destination으로 이어지는 길이 없을 수 있다는 뜻으로 해석할 수 있다.
접근 방법
강철 부대와 다른 지점 간 연결 정보가 필요하다.
연결 정보가 필요한 이유는, 강철 부대로부터 다른 노드를 방문하기 위해서는 서로 연결되어 있어야 하며, 이를 통해 두 지점 사이의 거리를 알 수 있기 때문이다.
문제의 input 중, 2차원 정수 배열을 roads로 담아서 주고, 이는 두 지점 간에 왕복이 가능하다는 뜻이므로, 이 배열을 순회하면서 인접 리스트를 완성해준다.
def solution(n, roads, sources, destination):
# 인접 리스트 만들기
adj = [[] for _ in range(n + 1)]
for a, b in roads:
# 서로 왕복할 수 있어야 하므로 각 인덱스 모두 목적지를 넣어준다.
adj[a].append(b)
adj[b].append(a)각 노드와 다른 노드 사이의 연결 정보를 얻었으면, 해당 정보를 이용하여 강철 부대의 위치로부터 다른 지점까지의 거리를 구한다.
복귀하는 경로를 구하는 것은 곧 강철부대로부터 해당 지점까지의 거리와 동일하다는 뜻이기 때문에, 그리고 모든 지점에서 강철 부대로 복귀하는 거리를 구하는 것이 문제이므로, 이는 강철 부대를 시작점으로 하고 해당 지점으로부터 각 지점까지의 거리를 구하라는 뜻과 같다.
거리를 구하기 위하여 크게 3가지 방법을 사용할 수 있다.
- dfs
- bfs
- dijkstra(heap 사용)
우선 dfs와 bfs는 두 노드 사이의 간선의 가중치가 1인 경우에만 이용할 수 있는 방법이다. 이 문제에서는 두 지역간의 길을 통과하는데 걸리는 시간은 모두 1로 동일하기 때문에, 두 방법 모두 이용할 수 있다.
다만, 문제에서 노드의 개수가 최대 10만개이기 때문에, 재귀적 방법인 dfs를 이용한다면 파이썬의 경우 stackoverflow error가 발생할 수 있다. 파이썬의 허용된 재귀 깊이는 1000이기 때문이다. 따라서 해당 문제를 dfs 방식으로 해결하기 위해서는 stack 제한을 풀어야 한다.
import sys
sys.setrecursionlimit(100000)굳이 dfs로 해결해야 할 필요는 없기 때문에, bfs 또는 dijkstra로 해결하는 것이 좋은 선택지일 수도 있다.
해당 문제는 bfs, dijkstra 모두 답을 도출해낼 수 있다. bfs의 경우 시간복잡도는 O(V + E) = O(10만 + 50만)이고, dijkstra의 경우 시간복잡도는 O((V + E)logV)이다. 이번에 푸는 방식은 다익스트라 알고리즘을 이용해보기로 한다.
다익스트라 알고리즘을 통해, 강철 부대로부터 각 지점까지의 위치를 구한다.
다익스트라 알고리즘의 자세한 원리는 생략한다.
- 우선 distance 배열을 INF로 초기화한다. 최댓값으로 초기화 하는 이유는 최솟값을 구해야 하기 때문이다.
- 우선 시작점을 heap에 넣은 뒤, 노드를 꺼내면서 수행한다. 이때
distance[next]가distance[node] + 1보다 작다면, 더 이상 탐색을 수행하지 않는다. distnace[next] > distnace[node] + 1이라면,distance[next]를 갱신하고 큐에 삽입한다.
def dijkstra(start, adj, distance):
q = []
distance[start] = 0
heapq.heappush(q, (0, start))
while q:
dist, node = heapq.heappop(q)
# 메모된 거리가 dist보다 짧다면 더 이상 탐색할 필요가 없다.
if distance[node] < dist:
continue
for nxt in adj[node]:
nxt_distance = dist + 1
if distance[nxt] > nxt_distance:
distance[nxt] = nxt_distance
heapq.heappush(q, (nxt_distance, nxt))
source를 바탕으로 정답을 반환한다.
dijkstra를 성공적으로 실행했다면, distance 배열은 강철 부대로부터 거리가 저장되어 있을 것이다. 이를 source 순서에 맞게 반환하면 된다. 이 때 주의해야 할 점은, 강철 부대로 부터 길이 존재하지 않아 INF의 값을 가진다면 -1을 출력할 수 있도록 바꿔주어야 한다.
정답 코드
import heapq
INF = 987654321
def dijkstra(start, adj, distance):
q = []
distance[start] = 0
heapq.heappush(q, (0, start))
while q:
dist, node = heapq.heappop(q)
# 메모된 거리가 dist보다 짧다면 더 이상 탐색할 필요가 없다.
if distance[node] < dist:
continue
for nxt in adj[node]:
nxt_distance = dist + 1
if distance[nxt] > nxt_distance:
distance[nxt] = nxt_distance
heapq.heappush(q, (nxt_distance, nxt))
def solution(n, roads, sources, destination):
adj = [[] for _ in range(n + 1)]
distance = [INF] * (n + 1)
answer = []
# create graph
for a, b in roads:
adj[a].append(b)
adj[b].append(a)
# get distance array
dijkstra(destination, adj, distance)
# return value by sources
for source in sources:
value = distance[source]
answer.append(value if value < INF else -1)
return answer배운 것
당연히 다익스트라 알고리즘이 더 빠를 줄 알았는데, 시간 복잡도를 찾아보니까 bfs가 선형적인 시간복잡도를 가지기 때문에 조금 더 빨랐다. 간선의 가중치가 1인 경우에는 bfs를 사용하는 것이 좋은 방식일 수도 있을 듯 하다.
