Python + scipy を使って最短経路もついでに取得する方法

こんばんは。kawap23です。

前回の記事で、最短経路知りたいなぁって言っているとTwitterで教えてくださった方がいるので、記事にして残しておきたいと思います。(maspyさん、ありがとうございます。)

あくまで、使うためですので、正確性や厳密性は保証しません。

kawap23.hatenablog.com

scipyを使って最短距離問題を解く方法にいついては、前回の記事を見てください。

最短経路を求める方法

いきなり結論ですが、dijsktra(グラフ、有向 or 無向、始点、経路が必要 or 不要)と引数に一つ足すだけで、最短距離と同時に返ってきます。(numpy.ndarrayが2つ返ってくる)

経路：return_predecessors = True で経路が返ってくる (デフォルトはFalse)
返ってきたarrayには直前の頂点情報が入っているため、始点までたどることで最短経路をたどることができる。
始点には -9999 が入っている
始点と連結していない点にも -9999が入る

とこんな感じです。ちなみに始点に2つ以上リスト等の形で指定しても問題なく返ってきます。

使用例は以下の感じです。入力例は蟻本のものに非連結成分を追加したものです。

f:id:kawap23:20190909224151p:plain — グラフ

import numpy as np
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import dijkstra

N, M = map(int, input().split())
edge = np.array([input().split() for _ in range(M)], dtype = np.int64).T
graph = csr_matrix((edge[2], (edge[:2] - 1)), (N, N))


distance_mat, processors = dijkstra(graph, directed = False, indices = [0, 1], return_predecessors = True)
print ('点0からの最短距離')
print (distance_mat)

print ('各点の直前の点')
print (processors)

# 入力例
# 9 11
# 1 2 2
# 1 3 5
# 2 3 4
# 2 4 6
# 2 5 10
# 3 4 2
# 4 6 1
# 5 6 3
# 5 7 5
# 6 7 9
# 8 9 4

# 出力例
# 点0からの最短距離
# [[ 0.  2.  5.  7. 11.  8. 16. inf inf]
#  [ 2.  0.  4.  6. 10.  7. 15. inf inf]]
# 各点の直前の点
# [[-9999     0     0     2     5     3     4 -9999 -9999]
#  [    1 -9999     1     1     1     3     4 -9999 -9999]]