[알고리즘] 프로그래머스 부대복귀 - Python

문제 : 프로그래머스 부대복귀 문제

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

프로그래머스

 

이 문제는 기본적인 그래프 dfs/bfs 문제이다. 구현을 하는 과정에서 bfs를 사용하였고, 그래프는 answer 배열을 활용하여 표현해보았다. 

from collections import deque
def solution(n, roads, sources, destination):
    answer = [[]*(n+1) for _ in range(n+1)]
    
    for i in range(len(roads)):
        answer[roads[i][0]].append(roads[i][1])
        answer[roads[i][1]].append(roads[i][0])
    
    def bfs(location):
        q=deque()
        q.append(location)
        visit=[0]*(n+1)
        visit[location]=1
        while q:
            cur_loc=q[0]
            q.popleft()
            if len(answer[cur_loc])>0:
                for i in answer[cur_loc]:
                    if visit[i]==0:
                        visit[i]=visit[cur_loc]+1
                        if i==destination:
                            return visit[i]-1
                        q.append(i)
        
        return -1
    temp=[]
    for i in sources:
        if i==destination: 
            temp.append(0)
        else: temp.append(bfs(i))
        
        
    return temp

 

bfs의 경우에는 따로 함수를 만들어서 사용하였다. location을 입력받으면 그 위치에서부터 그래프를 bfs로 탐색해나가는 것이다. visit 배열은 먼저 0으로 n+1의 크기만큼 초기화를 해둔다. 이 배열의 역할은 크게 두가지이다.

• 노드에 방문했는지 여부 확인 (0이면 방문하지 않은 노드, 0 초과이면 방문한 노드)
• 시작 노드로부터 해당 노드까지 거리 (시작점을 방문하고 1로 초기화를 했기 때문에 이 값에 -1을 한게 거리가 된다.)

이 코드에서는 시간초과가 발생한다. 그 이유는 bfs를 sources 배열의 원소 개수만큼 하기 때문이다. 단순하게 sources 배열에 있는 원소 개수만큼 매번 bfs를 돌리는 것이 아니라 한번 bfs를 하고 그 결과를 sources 배열을 순회하면서 사용하면 시간복잡도 측면에서 효율이 올라간다는 것을 알았다. 

 

또, 이 코드에서는 추가로 개선해야하는 점이 있다.

시작점을 sources의 원소, 끝점을 destination으로 지정하여 bfs안에서 destination을 만나면 종료를 하도록 해두었다. 이 부분도 만약, bfs 1회로 수정을 하게된다면 개선해야 하는 부분이었다. 그리고 이 코드에서는 visit 배열에서 해당 노드를 방문했는지 여부를 판단하기 위해 시작점을 1로 설정해두었다. (0이면 방문하지 않은것으로 설정해두었기 때문) 다만, 이렇게 하게 되면 리턴할 때 -1을 해야해서 번거롭다. 이 부분도 개선이 필요하다고 생각했다. 

 

이 개선점들을 반영하여 수정한 코드는 다음과 같다. 

from collections import deque
def solution(n, roads, sources, destination):
    answer = [[]*(n+1) for _ in range(n+1)]
    
    for i in range(len(roads)):
        answer[roads[i][0]].append(roads[i][1])
        answer[roads[i][1]].append(roads[i][0])
    visit=[0]*(n+1)
    def bfs(location):
        q=deque()
        q.append(location)
        
        visit[location]=1
        while q:
            cur_loc=q[0]
            q.popleft()
            if len(answer[cur_loc])>0:
                for i in answer[cur_loc]:
                    if visit[i]==0:
                        visit[i]=visit[cur_loc]+1
                        q.append(i)
        return -1
    temp=[]
    bfs(destination)
    
    for i in sources:
        temp.append(visit[i]-1)
    
    return temp

 

코드도 훨씬 간결해졌다. 총 1회의 bfs만으로 destination으로부터 모든 노드까지의 거리를 구하고 바로 답을 구할 수 있기 때문이다. 또한, 만약 도달하지 못한 경우는 -1을 리턴하도록 해야하는데 이 부분도 도달하지 못한 경우에는 0이 노드 값이기 때문에 -1을 해서 표현할 수 있었다. 

 

시간 복잡도는 bfs를 1회하면서 O(n)으로 개선이 되었다.