[프로그래머스] 배달

출처: 프로그래머스 코딩 테스트 연습, [프로그래머스] 배달

새로 알게된 것들

import heapq : 우선순위 큐를 사용하기 위한 힙 모듈
heapq.heappush(heap, item) : item을 heap에 추가
heapq.heappop(heap) : heap에서 가장 작은 원소를 pop & 리턴. 비어 있는 경우 IndexError가 호출됨.
일반적인 힙의 구조에서는 리스트의 첫번째 원소가 가장 작은 원소로 알고 있어 heap[0]을 이용해서 가장 작은 원소를 조회했는데 힙 리스트의 첫번째에 가장 작은 원소가 들어있지 않았음
heapq.heapify(x) : 리스트 x를 즉각적으로 heap으로 변환함 (in linear time, O(N))

풀이

한 마을에서 부터 나머지 마을까지 최단거리를 구해야하므로 다익스트라를 떠올릴 수 있었다

1. 마을과 거리에 대한 리스트를 사용하기 쉬운 형태로 변환
2. 각 마을의 방문여부 리스트만 만들고 각 마을의 거리에 대한 정보를 담는 리스트는 만들지 않음
3. 시작점 마을 1을 우선순위 큐에 넣고 반복문 시작
4. 이미 방문한 마을인 경우 스킵
5. 다익스트라 알고리즘 특성상 이전보다 더 작은 거리인 마을이 pop 될 수 없기 때문에 pop한 마을의 거리가 K보다 클 경우 반복문을 빠져나오고 카운팅한 방문을 반환
6. 4,5번에 해당되지 않을 경우 현재 마을에서 갈 수 있는 마을들을 우선순위큐에 누적거리와 함께 push
7. 위의 과정 반복

코드

import heapq

def solution(N, road, K):
    # {시작점: [(거리, 도착점), ...]}
    map = {i:[] for i in range(1, N + 1)}
    visited = [False for _ in range(N + 1)]
    for v in road:
        map[v[0]].append((v[2], v[1]))
        map[v[1]].append((v[2], v[0]))

    # [[누적거리, 도착점],...]
    heap = [[0, 1],]
    visited_count = 0
    while heap and visited_count != N:
        total, curr = heapq.heappop(heap)
        if visited[curr]: continue
        if total > K: break
        visited[curr] = True
        visited_count += 1
        for dist, next in map[curr]:
            if not visited[next]:
                heapq.heappush(heap, [total + dist, next])
    
    return visited_count

느낀점

저번 학기 알고리즘 수업을 듣고 다익스트라 알고리즘이 익숙해 조금 수월하게 풀었다. 다른 사람들도 우선순위큐를 이용해서 가장 작은 거리를 구했을 거라고 예상하며 다른사람 풀이는 보는데 내 코드와 많이 달라 놀랐다. 문제를 푸는 동안 어떻게 하면 더 시간을 줄일 수 있을까라는 생각을 했는데 일단 문제는 풀고 그런 생각을 해야겠다. 그리고 구조적으로 크게 시간을 줄이지 못하면 코드를 더 간결하고 더 단순하게 짜는 연습이 필요해 보인다.