【Atcoder】132E/Dijkstra法

今日の問題はこちら

先日勉強したDijkstra法の発展です。
今回は有向グラフかつ、またひとつの頂点に3回まで(3の剰余の分だけ)止まってもと良いというように問題を言い換えることができます。ひとつの頂点に3種類の経路があり得るので、経路はd[ノード数][3]のベクトルとします。

入力例1の場合、頂点3に至るコストは一週目は2(→d[頂点3][2%3]に格納)、二週目は6(→d[頂点3][6%3]に格納)となり、dを更新して行きます。

答えはd[ゴールの頂点][0]　(→ゴールの頂点までちょうど3の倍数まで行け、その場合のコストを表す)を3で割って完了！！

##解法

from heapq import *
f_inf = float('inf')

def dijkstra(start,node,edge):
    #始点sから各頂点への距離
    #node:頂点数
    #d:各頂点へのコスト(存在しないときはinf)
    #q:キュー(スタートからのコスト,頂点)
    d = [[f_inf, f_inf, f_inf] for _ in [0]*node]
    d[start][0] = 0
    #スタートをキューに入れる
    q = [(0,start)]

    while(len(q) != 0):
    #キューの先頭を取得
      ci,i = heappop(q)
    #キューの先頭から繋がっている頂点を探索
      for j in edge[i]:
        nc = ci + 1
        if (d[j][nc%3] > nc):
          d[j][nc%3] = nc
          heappush(q,(nc, j))
    return d
#入力
N,M = (int(x) for x in input().split())
#エッジ
E = [[] for _ in [0]*N]

for i in range(M):
  a,b = (int(x) for x in input().split())
  E[a-1].append(b-1)

S,T = (int(x)-1 for x in input().split())

d = dijkstra(S, N, E)
ans = d[T][0]

#出力
print(ans//3 if ans != f_inf else -1)

グラフ系はやっぱり苦手意識がすごい。。。
各頂点$x$回通る可能性がある場合、ノード数を$x$倍してしまえばいいんですね。
次こそ解けるようにしたい！