여러 부대원이 서로 다른 지역에서 강철부대가 있는 목적지로 복귀하려고 한다.
모든 길은 왕복할 수 있고, 길 하나를 지나는 시간은 모두 1이다.
각 출발 지역에서 목적지까지 복귀하는 최단시간을 구해야 한다.
목적지까지 도달할 수 없는 지역은 -1을 반환한다.
각 부대원의 출발 지역마다 BFS를 수행하면 sources의 원소가 많을 때 같은 탐색을 반복하게 된다.
모든 부대원의 목적지는 destination으로 같다.
따라서 목적지에서 시작해 한 번만 BFS를 수행하면 모든 지역에서 목적지까지의 최단거리를 동시에 구할 수 있다.
길은 양방향이므로 다음 두 거리는 같다.
지역 A -> 목적지
목적지 -> 지역 A
목적지에서 BFS를 수행해 구한 거리를 그대로 각 부대원의 복귀 최단시간으로 사용할 수 있다.
모든 길의 이동 시간이 1로 동일하다.
가중치가 모두 같은 그래프에서 BFS는 시작점으로부터 각 정점까지의 최단거리 순서대로 탐색한다.
목적지에서 가까운 지역부터 차례대로 방문하므로, 처음 방문한 순간의 거리가 최단시간이다.
거리 0: destination
거리 1: destination과 직접 연결된 지역
거리 2: 거리 1 지역과 연결된 아직 방문하지 않은 지역
...
각 길은 왕복할 수 있으므로 양쪽 지역의 인접 리스트에 모두 추가한다.
graph = [[] for _ in range(n + 1)]
for start, end in roads:
graph[start].append(end)
graph[end].append(start)
지역 번호는 1부터 n까지 사용하므로 크기가 n + 1인 배열을 사용한다.
각 지역에서 목적지까지의 최단시간을 저장할 배열을 만든다.
아직 방문하지 못한 지역은 -1로 둔다.
distance = [-1] * (n + 1)
목적지에서 목적지까지의 거리는 0이다.
distance[destination] = 0
목적지를 큐에 넣고 BFS를 시작한다.
queue = deque([destination])
큐에서 현재 지역을 꺼낸 뒤, 연결된 지역을 확인한다.
current = queue.popleft()
for neighbor in graph[current]:
아직 방문하지 않은 지역이라면 현재 거리보다 1 큰 값을 저장하고 큐에 추가한다.
if distance[neighbor] == -1:
distance[neighbor] = distance[current] + 1
queue.append(neighbor)
이미 방문한 지역은 더 짧거나 같은 경로로 탐색된 상태이므로 다시 확인하지 않는다.
모든 지역의 거리 계산이 끝난 뒤 sources의 순서대로 거리 배열 값을 꺼낸다.
return [distance[source] for source in sources]
도달할 수 없는 지역의 거리는 처음 값인 -1로 남아 있으므로 별도의 예외 처리 없이 요구사항을 만족한다.
from collections import deque
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)
# 각 지역에서 destination까지의 최단시간
distance = [-1] * (n + 1)
distance[destination] = 0
queue = deque([destination])
# 목적지에서 시작하는 BFS
while queue:
current = queue.popleft()
for neighbor in graph[current]:
# 처음 방문한 경로가 최단 경로다.
if distance[neighbor] != -1:
continue
distance[neighbor] = distance[current] + 1
queue.append(neighbor)
return [distance[source] for source in sources]
graph = [[] for _ in range(n + 1)]
각 지역과 직접 연결된 지역들을 저장하는 인접 리스트다.
인접 행렬보다 필요한 메모리가 적고, 현재 지역과 연결된 길만 확인할 수 있어 BFS에 적합하다.
distance = [-1] * (n + 1)
각 지역에서 목적지까지의 최단시간을 저장한다.
값이 -1이면 BFS로도 방문하지 못한 지역이므로 목적지까지 복귀할 수 없다는 뜻이다.
queue = deque([destination])
distance[destination] = 0
각 source에서 목적지로 가는 경로를 따로 찾지 않는다.
목적지에서 모든 지역으로 퍼져 나가면 한 번의 BFS로 모든 source의 답을 구할 수 있다.
if distance[neighbor] != -1:
continue
거리가 이미 저장된 지역은 방문한 상태다.
BFS는 거리가 가까운 지역부터 방문하므로 첫 번째로 저장한 거리가 최단거리다.
다음과 같은 길 정보가 있다고 하자.
roads = [
[1, 2],
[2, 3],
[3, 4],
[2, 5]
]
destination = 2
sources = [1, 3, 4, 5, 6]
2번 지역에서 BFS를 수행하면 거리 배열은 다음과 같다.
1번 지역: 1
2번 지역: 0
3번 지역: 1
4번 지역: 2
5번 지역: 1
6번 지역: -1
따라서 sources 순서에 맞춘 결과는 다음과 같다.
[1, 1, 2, 1, -1]
지역의 수를 N, 길의 수를 R이라고 하자.
인접 리스트를 만드는 데 O(R), BFS에서 각 지역과 길을 한 번씩 확인하는 데 O(N + R)이 필요하다.
O(N + R)
sources의 결과를 만드는 데는 O(S)가 추가되며, S는 sources의 길이다.
인접 리스트, 거리 배열, BFS 큐를 사용한다.
O(N + R)
이 문제는 여러 출발점에서 하나의 목적지로 가는 최단거리를 구하는 BFS 문제다.
모든 길을 양방향 인접 리스트로 구성
destination에서 BFS 시작
각 지역의 최단 복귀 시간을 distance에 저장
sources 순서대로 distance 값 반환
출발점마다 BFS를 반복하지 않고 목적지에서 한 번만 BFS를 수행하는 것이 핵심이다.