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

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

프로그래머스

각 시작지점 에서 목표지점 까지의 최단거리를 구하는 문제이다.

처음풀이는 문제요구사항 그대로 각각의 시작지점에서 목표지점까지 BFS로 거리를 구하였다.

하지만 11, 12, 13, 14, 15번 테케에서 시간초과..

from collections import defaultdict
from collections import deque


def solution(n, roads, sources, destination):
    def bfs(curr):
        visited = [False] * (n + 1)
        queue = deque()
        queue.append([curr, 0])
        visited[curr] = True
        while queue:
            x, count = queue.popleft()
            if x == destination:
                result.append(count)
                return
            for v in graph[x]:
                if not visited[v]:
                    visited[v] = True
                    queue.append([v, count + 1])
        result.append(-1)

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

    result = []
    for source in sources:
        bfs(source)
        
    return result


print(solution(3, [[1, 2], [2, 3]], [2, 3], 1))
print(solution(5, [[1, 2], [1, 4], [2, 4], [2, 5], [4, 5]], [1, 3, 5], 5))

반대로 생각해보면 목표지점에서 각 시작지점 까지의 최단 거리를 구하면 된다.

이럴때 사용하는 것이 다익스트라 알고리즘

from collections import defaultdict
from collections import deque


def solution(n, roads, sources, destination):
    def bfs(target):
        q = deque()
        q.append([0, target])
        dp[destination] = 0
        while q:
            cost, curr = q.popleft()
            for v in graph[curr]:
                if dp[v] == int(1e9):
                    q.append([cost + 1, v])
                    dp[v] = cost + 1

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

    dp = [int(1e9)] * (n + 1)
    bfs(destination)
    result = [dp[x] if dp[x] != int(1e9) else -1 for x in sources]
    return result

모든 거리의 비용이 1이기 때문에 priority queue를 사용하지 않았다.

 

목표지점에서 다른 지점까지의 거리를 무한대(int(1e9))로 설정한다.

처음에 목표지점을 queue에 넣은 뒤 거리를 0으로 설정한다.

각각의 연결된 노드를 탐색하면서 방문하지않은 노드만 cost를 갱신한다.

모든 노드간 간선의 가중치가 1이기 때문에 노드를 방문한 순간의 비용이 최소비용이다.

탐색후 거리가 업데이트 되지않은 노드는 방문 할 수 없는 노드이다.