ダイクストラ法の考え方

はじめに

初めてダイクストラ法を用いる問題と出会ったため備忘録として書きます。
https://atcoder.jp/contests/abc340/tasks/abc340_d
この問題です。

ダイクストラ法を使うとき

無向グラフや有向グラフにて任意の点から点までの最短距離を求める際にダイクストラ法を用いることができます。

ダイクストラ法の考え方

エッジに重みがなく、純粋に最短経路を求める際には幅優先探索と同じように記述します。エッジに重みがある場合は少し複雑になり、ヒープを使う必要があります。ヒープで出発点からの距離が最小になるような値を取り出していき、到達したノードから伸びているエッジをヒープに入れるというサイクルを繰り返していき、目的地がヒープから取り出されたらその値が最短距離となります。

コード例

https://atcoder.jp/contests/abc340/tasks/abc340_d
この問題のコード例です。TLEを出してしまったコードとACとなったコードの両方を掲載します。

誤った例.py

from collections import deque
import heapq
import sys

route_list={}


appended_node={1:0}

N=int(input())


for i in range(1,N):
    route_list[i]=[]
    a,b,x=map(int,input().split())
    route_list[i].append((a,i+1))
    route_list[i].append((b,x))

h=[]

for k in route_list[1]:
    heapq.heappush(h,k)




while h:
    sum,node=heapq.heappop(h)
    if node==N:
        print(sum)
        sys.exit()
    else:
        appended_node[node]=sum
        for k in route_list[node]:
            if k[1] not in appended_node:
                heapq.heappush(h,(k[0]+sum,k[1]))

このコードでは、未到達であったノードから伸びているエッジを無条件でヒープに入れました。そのため、ヒープに値が貯まりやすいです。

AC.py

from collections import deque
import heapq
import sys

route_list={}
# {2:[(100,3),(50,4)]}



N=int(input())

INT=10**19

dist=[INT]*(N+1)

for i in range(1,N):
    route_list[i]=[]
    a,b,x=map(int,input().split())
    route_list[i].append((a,i+1))
    route_list[i].append((b,x))

h=[]

for k in route_list[1]:
    heapq.heappush(h,k)




while h:
    sum,node=heapq.heappop(h)
    if dist[node]<sum:
        continue
    if node==N:
        print(sum)
        sys.exit()
    else:
        for k in route_list[node]:
            if dist[k[1]]>(k[0]+sum):
                dist[k[1]]=(k[0]+sum)
                heapq.heappush(h,(k[0]+sum,k[1]))

このコードではヒープに入れる値が最小値となりうるか事前にチェックしているため、計算量が安定している。