[프로그래머스] LV.3 부대복귀 / 파이썬(Python)

📜 문제 설명

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

 

🔗 문제 링크

https://school.programmers.co.kr/learn/courses/30/lessons/132266

프로그래머스

 

💡 문제 풀이

전형적인 BFS문제. 라고 생각했으나 전형적으로 BFS를 구현했다가 시간초과가 처참했다.

처음에 생각했던 방식은 출발점으로부터 목적지까지 걸리는 거리를 계산 + 예외 케이스에 대한 처리였는데 시간 초과랑 실패가 처참한 결과가 나왔다. 테스트 케이스가 최대 50만, 10만 까지 커질 수 있어서 최악의 경우에는 500만 번을 돌기 때문에 시간 초과가 처참했다. 

그래서 두번째로 시도한 방식이 목적지로부터 반대로 출발점까지 걸리는 시간을 계산하는 것이다.

그리고 걸리는 시간을 알아내야하는 출발지들만 리턴한다.

통과했다.

# 실패 코드
def solution(n, roads, sources, destination):
    answer = []
    
    V = [[] for _ in range(n + 1)]
    
    for i in range(len(roads)):
        V[roads[i][0]].append(roads[i][1])
        V[roads[i][1]].append(roads[i][0])
        
    for i in range(len(sources)):
        visited = [0 for _ in range(n + 1)]
        
        q = []
        q.append(sources[i])
        visited[sources[i]] = 1
        cnt = 0
        
        if sources[i] == destination:
            answer.append(cnt)
            break
        
        while q:
            cnt += 1
            t = q.pop(0)
            
            if 0 not in visited and cnt >= len(visited):
                cnt = -1
                answer.append(cnt)
                break
                
            if destination in V[t]:
                answer.append(cnt)
                break
            else:
                if len(V[t]):
                    for j in V[t]:
                        if not visited[j]:
                            visited[j] = 1
                            q.append(j) 
                else:
                    cnt = -1
                    answer.append(cnt)
                    break
                  
    return answer

 

def solution(n, roads, sources, destination):
    answer = []
    
    # 그래프 만들기
    V = [[] for _ in range(n + 1)]
    
    for s, e in roads:
        V[s].append(e)
        V[e].append(s)
    
    # 방문 체크 -> 도달할 수 없는 경우 리턴값은 -1이므로 디폴트값으로 설정
    visited = [-1 for _ in range(n + 1)]
    visited[destination] = 0
    
    # 목적지로 부터 출발해서
    q = [destination]
    while q:
        t = q.pop(0)
        
        # 그래프가 이어진 길을 따라가면서 걸린 시간을 1씩 더한다
        for i in V[t]:
            if visited[i] == -1:
                q.append(i)
                visited[i] = visited[t] + 1
    
    # 거리를 계산해야 하는 출발지만 리턴
    return [visited[i] for i in sources]