【Python】有向グラフの最短経路に含まれる辺を列挙

辺の長さが1の有効グラフの最短経路に含まれる辺を列挙する方法メモ
最短経路が複数存在する場合は、そのうちの1つを出力する

from collections import deque,defaultdict
# 頂点数:n
# 辺の数:m
n,m=map(int,input().split())
edge=defaultdict(list)
for i in range(m):
    # 0-indexにする
    s,t=map(lambda x:int(x)-1,input().split())
    # 遷移先の頂点と辺のindexを追加する
    edge[s].append((t,i))

INF=float('inf')
# 頂点0からの距離
dist=[]
# ひとつ前の頂点と辺のindex
back=[(-1,-1)]*n

def bfs(s):
    global dist,back
    dist=[INF]*n
    dist[s]=0

    q=deque()
    q.appendleft(s)
    while len(q)>0:
        now=q.pop()
        for to ,ind in edge[now]:
            if dist[to]!=INF: continue
            dist[to]=dist[now]+1
            back[to]=(now,ind)
            q.appendleft(to)

bfs(0)

used=set()
now=n-1
while now!=0:
    before, ind =back[now]
    used.add(ind)
    now=before

print(used)

実行結果

入力

4 4
1 2
2 3
2 4
3 4

出力

{0, 2}

参考