[Python/BFS] 부대 복귀

문제 설명

강철부대의 각 부대원이 여러 지역에 뿔뿔이 흩어져 특수 임무를 수행 중입니다. 지도에서 강철부대가 위치한 지역을 포함한 각 지역은 유일한 번호로 구분되며, 두 지역 간의 길을 통과하는 데 걸리는 시간은 모두 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	source	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합니다.

---

풀이

첫 접근법에는 각 source에서 출발해서 destination까지 가는 거리를 계산하고자 했다.

하지만 여러 번 탐색을 해야해서 비효율적이고 dest에 도착하지 못하는 경우를 따로 처리해줘야 한다 (아래 코드는 처리도 안함)

# *** 완전 오답 ***
from collections import deque
def bfs(n, visited, nodes, start, end):
    q = deque()
    q.append(start)
    visited[start] = 1
    num = 0
    while q:
        tmp = q.popleft()
        num += 1
        if tmp == end:
            return num
        for i in nodes[tmp]:
            if visited[i] == 0:
                q.append(i)
                visited[i] = 1       
        
def solution(n, roads, sources, destination):
    answer = []
    nodes = [[] for _ in range(n+1)] # 인접 리스트
    for road in roads:
        a, b = road[0], road[1]
        nodes[a].append(b)
        nodes[b].append(a)
    # print(nodes)
    
    visited = [0]*(n+1)
    for source in sources:
        answer.append(bfs(n, visited, nodes, source, destination))
    return answer

 

세 가지 점을 유의해야 한다.

1. destination에서 출발
2. 결과 list를 -1로 선언해두고 +1씩
3. visited 기록을 위해서 list가 아닌 set 사용

'''
[input]
n: 지역 수
roades: 두 지역을 왕복할 수 있는 길의 정보 
sources: 각 부대원이 위치한 서로 다른 지역을 나타내는 정수
destination: 강철부대의 지역

[output]
sources의 원소 순서대로 강철부대로 복귀할 수 있는 최단시간을 담은 배열

[idea]
BFS - deque, 인접 리스트, visited
'''
from collections import deque    
        
def solution(n, roads, sources, destination):
    answer = []
    graph = [[] for _ in range(n+1)] # 인접 리스트
    for road in roads:
        a, b = road[0], road[1]
        graph[a].append(b)
        graph[b].append(a)
    # print(nodes)
    
    ans = [-1 for _ in range(n+1)]
    ans[destination] = 0
    Q = deque([destination])
    visited = set([destination]) # 원소 중복 없이 add하려고 집합 사용
    while Q:
        node = Q.popleft()
        for x in graph[node]:
            if x not in visited:
                visited.add(x)
                ans[x] = ans[node]+1
                Q.append(x)
                
    return list(ans[s] for s in sources)

 

📝

인접 리스트 선언시 default dict 사용 가능

graph = defaultdict(list)
for a,b in roads:
    graph[a].append(b)
    graph[b].append(a)

 

 

참고