[Python Algorithm] 백준 최단 경로 - 1753

Chani·2022년 3월 1일
0

알고리즘

목록 보기
6/16
post-thumbnail


풀이 과정

다익스트라 알고리즘을 사용하여 구현하면 된다.
처음에는 힙큐에 데이터를 넣어줄때 (정점, 비용) 순서로 넣어주는게 자연스러울것 같아서 파이썬의 클래스와 스페셜 메소드 lt를 사용하여 힙큐의 우선순위를 커스텀 하는 방식으로 구현하였다.
코드는 다음과 같다.

class node:
  def __init__(self, v, w):
    self.v = v
    self.w = w

  def __lt__(self, other):
    return self.v < other.w # 오름차순
    
def dijkstra(start):
  heap = []
  distance = [INF] * (V + 1)

  distance[start] = 0
  heapq.heappush(heap, node(start, 0))

그러나 위와 같이 힙큐에 데이터를 넣어줄때마다 객체를 만들어서 넣어주게 되니 시간초과가 발생하였다.

결국 맘에는 안들지만 (비용, 정점) 순서대로 튜플 형태로 힙큐에 넣는 방식으로 구현하였고, 이렇게 구현하니 통과가 되었다.

비록 원하는 방향대로 구현하지는 못했지만 힙큐의 우선순위를 커스텀 해주는 방식은 이후 다른 알고리즘 문제를 해결할때 사용될 수 있다고 생각한다.


최종 코드

import sys
import heapq

input = sys.stdin.readline
INF = sys.maxsize

V,E = map(int, input().split(' '))
K = int(input())

graph = [[] for _ in range(V + 1)]

for i in range(E):
  u,v,w = map(int, input().split(' '))
  graph[u].append((v, w))

def dijkstra(start):
  heap = []
  distance = [INF] * (V + 1)

  distance[start] = 0
  heapq.heappush(heap, (0, start))

  while heap:
    weight, vertex = heapq.heappop(heap)

    if weight > distance[vertex]: continue

    for v, w in graph[vertex]:
       if distance[v] > weight + w:
         distance[v] = weight + w
         heapq.heappush(heap, (weight + w, v))

  for i in range(1, V+1):
    print(distance[i]) if distance[i] != INF else print('INF')

dijkstra(K)

결과


최단 경로 문제 출처
GitHub 코드

profile
프론트엔드에 스며드는 중 🌊

0개의 댓글