[백준] 1753. 최단 경로 python

🐳 링크

백준 1753번 최단경로

🐕 문제

방향그래프가 주어지면 주어진 시작점에서 다른 모든 정점으로의 최단 경로를 구하는 프로그램을 작성하시오. 단, 모든 간선의 가중치는 10 이하의 자연수이다.

첫째 줄에 정점의 개수 V와 간선의 개수 E가 주어진다. (1≤V≤20,000, 1≤E≤300,000) 모든 정점에는 1부터 V까지 번호가 매겨져 있다고 가정한다. 둘째 줄에는 시작 정점의 번호 K(1≤K≤V)가 주어진다. 셋째 줄부터 E개의 줄에 걸쳐 각 간선을 나타내는 세 개의 정수 (u, v, w)가 순서대로 주어진다. 이는 u에서 v로 가는 가중치 w인 간선이 존재한다는 뜻이다. u와 v는 서로 다르며 w는 10 이하의 자연수이다. 서로 다른 두 정점 사이에 여러 개의 간선이 존재할 수도 있음에 유의한다.

🐾 풀이 방법

처음에 풀이는, 인접행렬을 사용해서 dfs로 모든 경우를 다 돌면서 cost를 업데이트해주고, 가장 작은 cost를 result배열에 저장해주는 식으로 풀었다. 하지만 계속 메모리 초과가 났고, 결국 검색 끝에 dijkstra 및 heap, 그리고 인접행렬 대신 인접리스트를 사용해야한다는 것을 알게 되었다!!

먼저 입력을 받을 때 인접 리스트로 [index, cost]를 graph[start]에 넣어주고, dijkstra를 이용해서 풀이를 시작한다.
먼저 dist라는 배열을 선언한 뒤, dist[k]에는 0을 넣어준다.(조건)
heapq에 (이하 hq) [cost,node]를 넣어준 뒤, cost가 최소인 노드를 hq에서 pop하며 비교해준다. cost가 dist[node]의 값보다 크면 continue를 통해 다음으로 넘어가고, 아니라면 해당 node의 이웃을 탐색하며 cost를 업데이트한다.

또한 여기서, 주의할 점이 원래 항상 입력을 받고 분리시킬 때 map과 input을 활용했는데 이 부분에서 시간 초과가 뜰 줄은 상상도 못했다! 😅 sys.stdin.readline().split()으로 변경해주었더니 시간 초과를 해결할 수 있었다.. input이 더 느리다고 한다!

solution 코드

1. 인접리스트 및 dijkstra를 사용해서 성공한 코드

  import sys, heapq
  sys.stdin = open("input.txt","rt")
  inf = float('inf')

  def dijkstra(v,k,graph):  
  dist = \[inf for \_ in range(v+1)\]  
  dist\[k\] = 0

  hq = []
  heapq.heappush(hq, [0,k]) #heapq에 [cost,node]를 넣음

  while hq:
      temp = heapq.heappop(hq) #cost가 최소인 노드
      cost = temp[0]
      node = temp[1]

      if cost > dist[node]:
          continue

      for i in graph[node]:
          neighbor = i[0]
          cost_i = dist[node] + i[1]

          if cost_i < dist[neighbor]:
              dist[neighbor] = cost_i
              heapq.heappush(hq, [cost_i, neighbor])
  return dist


  v, e = \[int(x) for x in sys.stdin.readline().split()\]  
  k = int(input())  
  graph = \[\[\] for \_ in range(v)\]  
  for \_ in range(e):  
  start, end, weight = \[int(x) for x in sys.stdin.readline().split()\]  
  graph\[start\].append(\[end,weight\])

  answer = dijkstra(v,k,graph)

  for dis in answer\[1:\]:  
  print(dis if dis != inf else 'INF')

2. 인접행렬을 사용해서 메모리 초과가 났던 코드...

import sys
sys.setrecursionlimit(10**5)
sys.stdin = open("input.txt","rt")

def dfs(node, target, value):
    if node == target:
        if result[target] == -1:
            result[target] = value
        if result[target] > value:
            result[target] = value
    else:
        for i in range(v):
            if graph[node][i] != 0 and visit[node] == 0:
                visit[node] = 1
                dfs(i,target,value + graph[node][i])
                visit[node] = 0



if __name__ == "__main__":
    v, e = map(int,input().split())
    k = int(input())
    k -= 1

    graph = [[0] * v for _ in range(v)]
    for i in range(e):
        start, end, weight = map(int,input().split())
        graph[start-1][end-1] = weight

    result = [-1] * v
    visit = [0] * v
    for i in range(v):
        if i != k:
            dfs(k,i,0)
        else:
            result[i] = 0

    for res in result:
        if res == -1:
            print("INF")
        else:
            print(res)