SciPyによる経路探索

Pythonでの競技プログラミングにおいて、典型的なグラフ探索アルゴリズム（dijkstra, floyd_warshall等）を使う場合、使用可能ならSciPyに任せるのがよい。

scipyのFloyd-WarshallとDijkstraのすすめ　Python 競技プログラミング Atcoder - じゅっぴーダイアリー

グラフのデータ構造の作り方

scipy.sparse.csgraph には、グラフの作り方がいくつかあるが、csr_matrixを使うのが楽なんじゃないか、というメモ。

scipy.sparse.csr_matrix — SciPy v1.4.1 Reference Guide

一般に、グラフ問題の多くは「頂点 $a$ と頂点 $b$ をコスト $c$ で結ぶ辺」がずらーっと並んで入力が与えられることが多い。

  4 3       ※頂点数, 辺数
  1 2 10
  2 3 20    →    ①--10--②--20--③--30--④
  3 4 30

対して、頂点を $N$ 個として、 $N \times N$ 行列でそれぞれの距離を表すような入力であることもある。

  4
  -1 10 -1 -1
  10 -1 20 -1
  -1 20 -1 30
  -1 -1 30 -1  (辺の無いところは-1とする)

それぞれについて、読み込み方を書いておく。

辺による入力

先頭行に頂点数と辺数を示す N M が与えられ、その後 $M$ 行に渡って結ばれる2頂点とコスト A B C が与えられるとする。

N M
A1 B1 C1
：
AM BM CM

入力行数が多くなることもあるため、sys.stdin.buffer.read() などを使うと、input()をforループで回した場合と比べて高速に読み込める。

$A,B,C$ をそれぞれ縦に読んだ長さ $M$ のNumPy配列を作り、 csr_matrix((cost, (aaa, bbb))) とすると、グラフ探索アルゴリズムに与えるグラフを作成できる。

NumPyとかは与えるのがタプルでもリストでも柔軟に解釈してくれることもあるけど、この場合はリストだと多義的な解釈になってしまうためか、タプルである必要がある。

        
              import sys
import numpy as np
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import floyd_warshall
 
inp = list(map(int, sys.stdin.buffer.read().split()))
n, m = inp[:2]
si = 3 * m + 2
aaa = np.array(inp[2:si:3]) - 1  # 1～Nで頂点番号が付けられている場合、0～N-1にする
bbb = np.array(inp[3:si:3]) - 1
ccc = np.array(inp[4:si:3])
graph = csr_matrix((ccc, (aaa, bbb)), shape=(n, n))
 
dists = floyd_warshall(graph, directed=False).astype(np.int64)

隣接行列による入力

先頭行に頂点数を示す N が与えられ、その後 $N \times N$ 行の隣接行列が与えられるとする。

N
C11 C12 .. C1N
：      ..  ：
CN1 CN2 .. CNN

$N \times N$ のNumPy配列を作って csr_matrix(mat) と与える。

        
              import sys
import numpy as np
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import floyd_warshall
 
inp = list(map(int, sys.stdin.buffer.read().split()))
n = inp[0]
mat = np.array(inp[1:n * n + 1]).reshape((n, n))
mat = np.where(mat == -1, 0, mat)  # 辺の無い箇所が-1などで表される場合、0に置換
graph = csr_matrix(mat)
 
dists = floyd_warshall(graph, directed=False).astype(np.int64)

この方法の問題点としては、辺の無い箇所を“0”で示すという仕様上、コスト0の辺が含まれる場合に正しく機能しなくなってしまう点。

（辺による入力では、「コスト0の辺」として明示的に与えることで正しく機能する）

他の方法

主に以下の4つの方法があるが、生成したグラフ上で探索を行う場合はcsrでよい。

csr_matrix
- 説明済み
csc_matrix
- 内部処理的に行と列 (Row, Column) のどちらを取り出しやすいかが入れ替わっただけで、あまり変わらない
lil_matrix
- 隣接行列ベース。空の行列に値を1つ1つ設定していく場合は、これでまず作ってからcsrなどに変換した方がよい
coo_matrix
- 辺情報からグラフを高速に作るが、演算は出来ない（csrなどに変換する必要がある）