프로그래머스 배달 코드 및 해설 (파이썬)

2021. 7. 5. 14:57algorithm

반응형

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

 

코딩테스트 연습 - 배달

5 [[1,2,1],[2,3,3],[5,2,2],[1,4,2],[5,3,1],[5,4,2]] 3 4 6 [[1,2,1],[1,3,2],[2,3,2],[3,4,3],[3,5,2],[3,5,3],[5,6,1]] 4 4

programmers.co.kr

 

2가지 수정사항을 제외하고 다익스트라 알고리즘을 그대로 활용해 풀이했습니다. 

(1) 하나의 도로에서 양방향으로 이동 가능하므로, graph를 만들 때 양방향 정보를 모두 저장했습니다

(2) 마지막에 답을 구할 때 주어진 K 이하의 시간이 걸리는 마을의 개수를 세서 반환했습니다. 

 

import heapq

def solution(N, road, K):
    INF = int(1e9)
    graph = [[] for _ in range(N+1)]
    distance = [INF] * (N+1)
    
    # 간선 정보 저장하기 
    for r in road:
        a, b, c = r
        # 양방향
        graph[a].append((b,c))
        graph[b].append((a,c))
    
    # 다익스트라 알고리즘 
    def dijkstra(start):
        q = []
        distance[start] = 0
        heapq.heappush(q, (0, start))
        
        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]))
    
    dijkstra(start=1)
    
    # K 이하의 시간에 배달이 가능한 마을의 개수 
    return len([d for d in distance if d <= K])

 

반응형