[프로그래머스] 배달, Summer/Winter Coding(~2018), python

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

코딩테스트 연습 - 배달

from collections import deque
def solution(N, road, K):
    INF = 5000000
    g= [[] for _ in range(N + 1)]
    dist = [INF] * (N + 1)
    for i in range(len(road)):
        g[road[i][0]].append((road[i][1],road[i][2]))
        g[road[i][1]].append((road[i][0],road[i][2]))
    
    
    q = deque() # start = 1
    q.append((1, 0))
    dist[1] = 0
    while q:
        cur, weight = q.popleft()
        for i in range(len(g[cur])):
            if weight + g[cur][i][1] < dist[g[cur][i][0]]:
                dist[g[cur][i][0]] = weight + g[cur][i][1]
                q.append((g[cur][i][0], dist[g[cur][i][0]]))
    
    return len([i for i in dist if i <= K])

우선 순위 큐를 사용할까했지만 여러 방법을 사용해보고 싶어서 다익스트라 알고리즘을 큐로 구현했다.

원래 거리보다 갱신되는 거리가 더 짧을때 queue에 추가하는 방향으로 하면된다.

 

 

---

우선순위큐로 하려면 visit을 하나 잡고 거리로 우선순위를 잡아서 넣으면 된다.

from heapq import heappop, heappush
def solution(N, road, K):
    INF = 5000000
    answer = 0
    g= [[] for _ in range(N + 1)]
    for i in range(len(road)):
        g[road[i][0]].append((road[i][1],road[i][2]))
        g[road[i][1]].append((road[i][0],road[i][2]))
        
    dist = [INF] * (N + 1)
    visit = [0] * (N + 1)
    q = [(0, 1)]  # start = 1
    dist[1] = 0
    while q:
        weight ,cur = heappop(q)
        if visit[cur]: continue
        visit[cur] = 1
        for i in range(len(g[cur])):
            if weight + g[cur][i][1] < dist[g[cur][i][0]]:
                dist[g[cur][i][0]] = weight + g[cur][i][1]
                heappush(q, (dist[g[cur][i][0]], g[cur][i][0]))

    return len([i for i in dist if i <= K])

 

우선순위 큐가 조금더 빠르다.