[프로그래머스] Level 2 - 배달 (Python)

 

 

 

 

https://programmers.co.kr/learn/courses/30/lessons/12978

코딩테스트 연습 - 배달

 

 

 

 

✅ Solution

• 한 정점으로부터 다른 모든 정점까지의 최단 거리를 구해야 하므로 다익스트라 알고리즘을 이용한다.
• heapq를 이용하여 다익스트라 알고리즘을 구현한 뒤, K 이하의 거리를 가지는 정점의 수를 세어준다.
• 양방향, 즉 a->b, b->a 모두 가능하므로 주의하면서 간선 정보를 처리한다.

 

 

 

 

✅ Code

import heapq


def dijkstra(start, graph, distance):
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0

    while q:
        dist, now = heapq.heappop(q)
        if distance[now] < dist:
            continue

        for i in graph[now]:
            cost = dist + i[1]
            if cost < distance[i[0]]:
                distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))


def solution(N, roads, K):
    answer = 0
    INF = int(1e9)

    graph = [[] for _ in range(N + 1)]
    distances = [INF] * (N + 1)

    for road in roads:
        a, b, c = road
        graph[a].append((b, c))
        graph[b].append((a, c))

    dijkstra(1, graph, distances)

    for distance in distances:
        if distance <= K:
            answer += 1

    return answer