[프로그래머스] 부대 복귀

부대 복귀

Programmers

문제 설명

제한 사항

작성 코드

from collections import deque
def solution(n, roads, sources, destination):
    # 각 위치의 거리정보를 담아 두는 배열 (초기값은 -1)
    long = [-1]*(n+1)
    # 그래프를 dictionary로 선언.
    dic=dict()
    for i in range(1,n+1):
        dic[i]=[] #
    for i,j in roads:
        # road를 이용하여, 연결된 곳을 append
        dic[i].append(j)
        dic[j].append(i)
    q=[]
    q=deque() # BFS를 위한 deque() 선언
    visit=set() #방문한 곳을 집합으로 받음.
    q.append([destination,-1]) # Q 초기에 [목적지, -1]으로 설정(거꾸로 가는 형태)
    # BFS 시작
    while q:
        loc,num=q.popleft()
        #이미 방문한 적이있다면 넘어감
        if loc in visit:
            continue
        # 방문한 적이 없다면, 현재 위치를 방문표시하고, 최단거리를 기록함.
        visit.add(loc)
        long[loc]=num+1 #초기에 long[destination]에 0이 저장되도록
        # 도착한 지점에 연결되어있는 지역만큼 반복해서 append를 함.
        for i in dic[loc]:
            q.append([i,num+1])
    answer=[]
    # sources에 있는 배열만큼 정답을 채워넣음.
    for i in sources:
        answer.append(long[i])
    return answer