WUPC 2012 E - 会場への道 (python)

拡張ダイクストラの簡単な問題ということで、まず解いてみました。通常のダイクストラと何が違うかというと、例えば、頂点が3つなら最小距離を保存するリスト d も
d = [float("inf")] * 3 とかで十分でしたが、今回の問題のように、総距離が４の倍数になるものの中で最小のものを考える場合、リストを二次元に拡張して、
d[ i ][ j ] を“頂点 i に行くときの距離で、４で割ったときの余りが j の中で最小のもの”
としないと必要十分に答えを探せなくなる。
分かりやすい図はこちら
WUPC 2012：E - 会場への道 - 数学/競プロメモ
イメージとしては二次元 dp みたいな感じ。
まだ heap に入れるものが余分に１つ増えただけなのと、同じ頂点 d[ i ][ j ] は辿らなくてよいので簡単な方だと思われる。

import heapq
def dijkstra_heap(s,edge,p):
    d = [[float("inf")] * p for i in range(n)] 
# 通常のダイクストラと異なるところその１、あまりの数だけ頂点の数を増やす。（p倍する）
    used = [[True] * p for i in range(n)] # 今回は各頂点の最小値が更新されないので d の数だけ用意するばよい
    d[s][0] = 0
    used[s][0] = False
    ed_list = []
    for es in edge[s]:
        heapq.heappush(ed_list,[es[0],es[1],0]) 
# 異なるところその２、通常の（総距離、次に向かう頂点）に加えて、今の余りの情報が追加されている。
    while len(ed_list):
        poped = heapq.heappop(ed_list)
        md = poped[0]%p
        if not used[poped[1]][md]:
            continue
        v = poped[1]
        d[v][md] = poped[0]
        used[v][md] = False
        for e in edge[v]:
            if used[e[1]][(md+e[0])%p]:
                heapq.heappush(ed_list,[e[0]+poped[0],e[1],md])
        if not used[n-1][0]:
          break
    return d[n-1][0]

import sys
input = sys.stdin.readline
n,m = map(int,input().split())

edge = [[] for i in range(n)]
for i in range(m):
    x,y,z = map(int,input().split())
    if x != n-1:
      edge[x].append([z,y])
    if y != n-1:
      edge[y].append([z,x]) 
print(min(dijkstra_heap(0,edge,4),dijkstra_heap(0,edge,7)))