目次

SciPyによる経路探索

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

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

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

一般に、グラフ問題の多くは「頂点 $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でよい。