🚼

AtCoderのグラフ問題でscipy.sparseを使いたい

2023/02/19に公開

Python

AtCoder

scipy

tech

解説記事が少なくscipyの公式ドキュメントを読むようになってしまったので自分の経験をまとめ。
scipy.sparseはグラフ問題でよく使うアルゴリズムをいい感じに関数に丸投げして自前実装よりも高速で解いてくれるライブラリ。
scipy公式以外によくお世話になる貴重な日本語解説サイト。

入力受取

N頂点M個の辺でよくある入力形式。

N M
U1 V1
...
UM VM

import numpy as np
import sys
n, m = map(int, input().split())
u, v = np.fromstring(sys.stdin.read(), dtype = np.int64, sep = " ").reshape(-1,2).T-1

list内包表記等いろいろやり方はあるが、np.fromstringが楽と感じるようになった。sys.stdin.read()で入力の最後まで文字列として受け取り、sepで半角スペースを指定すると、改行コードも分割と認識してnumpy配列を作ってくれる。input()やsys.stdin.readline()では1行しか読み込んでくれず、大量の辺を一度に読み込めない。sys.stdin.readlines()はlistになるが、fromstring関数の第1引数はlistではなく文字列を渡す必要がある。sys.stdin.bufferとnumpy.frombufferはinput()と混ぜて使うと意図しない挙動になるので個人的には使うのを避けている。
reshapeは1次元配列を入力と同じM行2列に整形する。reshape中の-1はmとしても同じ。次のTは行列の転置を取るので、入力の辺の情報が2行M列に変形される。最後の-1は頂点番号を0始まりにするため。
4行目左辺のu, vがなんとなく違和感あるが、左辺にr個の変数を、右辺にr行c列の配列を持ってくると、各変数に各行の配列が割り当てられる。
実際の問題では他にも入力があることが多いので、最初に全ての入力をnumpy.fromstring関数で1つの配列にしてしまい、スライスで各変数に割り当てていく工夫がいる。入力にアルファベット等の文字列が混じったり、01の文字列がスペース無しで与えられる入力だと、fromstring関数が使えないので面倒。

グラフの作成

from scipy.sparse import coo_matrix
graph = coo_matrix((np.ones_like(u), (u, v)), shape = (n, n))

頂点、辺の情報を隣接行列として疎行列形式に変換する。importするのはcoo_matrixでもcsr_matrixでもcsc_matrixでもok。処理時間はcoo_matrixが気持ち早いが気持ちレベル。coo_matrixはcsr_matrix等に変換しないと行列計算ができないので、隣接行列のn乗を計算する必要がある場合はcsr_matrixで読み込むか.tocsr()メソッドでcoo_matrixからcsr_matrix形式に変換必要。グラフ問題では隣接行列の演算を自分で陽に計算することは少ないのでcoo_matrixでも十分。第一引数には(辺の重み, (頂点, 頂点))の二重括弧を入力する。辺の重みが関係なく、とにかく辺が張られていればよい場合は辺の重みとしてnumpy.ones_likeでも、頂点uをそのまま入れてしまってもよい。shapeは頂点数。疎行列は省メモリだが、頂点数が10^9で辺の数が10^5程度の場合10^9に相当するメモリが必要でAtCoderではメモリが足らない、計算が遅い、ということが起きるので、あらかじめ座標圧縮等で辺が張られていない頂点を削除する工夫が必要になる。
無向グラフの場合、u→vの辺に加えv→uの辺も加えたくなる（隣接行列の転置を足して対称行列にする）が、グラフ問題を解くscipyの関数は有向か無向か指定するオプションがあり、一方通行の辺情報をグラフに入力しても、解くときに無向を指定すれば勝手に逆向きの辺もあるとみなして計算してくれる。

グラフ問題をライブラリに頼る

使える関数

scipy.sparse.csgraph内にある関数を使う。AtCoderでお世話になるのはBFS（幅優先探索）、ワーシャルフロイド、connected componentsあたり。自分のレベルでは使いこなせていないが、最小全域木(minimum_spanning_tree)、最大フロー問題(maximum_flow)、二部グラフの最大マッチング(maximum_bipartite_matching)もおそらくライブラリ任せで高速に計算してくれると思う。

要注意（だった）関数（ダイクストラ法とDFS）

ダイクストラ法(dijkstra)とDFS（深さ優先探索 depth_first_order）は関数としては用意されているが、この2つの関数をAtCoderの問題で使うと数10あるテストケースのうち1つか2つでTLE（実行時間超過）になることが経験上多い。・・・と書いたが、2023年の言語アップデートで爆速になった。以前はDijkstra法は多数(>10^4個)の辺が接続したハブ頂点がある場合にTLEを起こしがちだったが、言語アップデート後はまったく問題なく通るようになった。たぶんDijkstraが想定解法でscipyのDijkstraが使えない、ということは無くなったと思う。。DFSの挙動は未確認。

ライブラリ丸投げの使い方

from scipy.sparse.csgraph import breadth_first_order
from scipy.sparse.csgraph import floyd_warshall
from scipy.sparse.csgraph import connected_components

# BFS幅優先探索 scipy内ではBFSではなくBFOという呼び方の関数
node, pred = breadth_first_order(graph, start, directed = False)
# ワーシャルフロイド scipy内（というか英語圏では？）逆順のfloyd_warshallという名前
dist = floyd_warshall(graph, directed = False).astype(np.int)
# connected components
n_components, label = connected_components(graph, directed = True, connection = "weak")

各関数にあるdirectedが有向か無向か指定するオプションで、デフォルトはTrue（有向グラフ）。BFSの最初の戻り値nodeは頂点startからBFSで頂点を探索したときの探索順のnumpy配列で、配列の最初の要素はstart。頂点startから辿り着ける頂点の数だけ要素がある。2番目の戻り値のpredはBFSで訪れる頂点の親の番号。n頂点のグラフの場合はn個の要素を持つ。startの親と、startから辿り着けない頂点の親は-9999となる。breadth_first_orderにはreturn_predecessorsという引数があり、デフォルトはTrueで第2戻り値にpredが返るようになっているが、いらない情報の場合はFalseにするとnodeだけが得られる。startからある頂点までの経路を求めるには、最後の頂点からpredを辿ってstartにたどり着くまでfor文かwhile文を回せば良い。が、numpy配列でfor文を回すと遅いので、経路復元でpredをfor文の中に入れる場合は、for文に入れる前にpred=pred.tolist()でnumpy配列からpython標準のlist形式に変換した方が良い。
ワーシャルフロイド法はi番目の頂点からj番目の頂点に向かう最短距離がi行j列に格納されたnumpy配列が得られる。整数を入力しても必ず浮動小数点型が返ってくるので、AtCoderの問題ではastypeで整数型に変換することが多い。入力するグラフは疎行列coo_matrixでなくてnumpy配列の行列形式でもよい。floyd_warshallはbreadth_first_orderと同じくreturn_predecessorsをオプションに持つがこちらはBFSと異なりreturn_predecessors=Falseがデフォルト。
connected_componentsの最初の戻り値n_componentsは連結成分数。2番目の戻り値labelは各頂点がどの連結成分に属するかのnumpy配列。有向グラフ(directed=True)の場合、connectionに"strong"を指定すると強連結成分分解となる。"weak"では頂点uと頂点vが一方通行でも辿り着ける場合同じ連結成分と判定される。

トポロジカルソート

いろいろ調べたがscipyでトポロジカルソートを一発で実現する方法は無さそう。が、2023年の言語アップデートでPython標準ライブラリで簡単に計算できるようになった。graphlibライブラリにあるTopologicalSorterをimportすることでトポロジカルソートを計算できる。AtCoderのトポロジカルソートの問題でいくつか使ってみた印象では、やや計算速度に不安があるが十分実用的。標準ライブラリなのでpypyならもっと早くなるかもしれない。また、AtCoderでは遅くて使えない、との悪評のあったnetworkxも言語アップデート後は普通に使えそうとの記事あり。

丸投げできない？

AtCoderの問題だとグラフを入力してscipyの関数に投げるだけだと解けなくてもう一工夫二工夫必要だったり、そもそも問題の見た目がグラフ問題に見えずグラフで解くという発想に至らないことがあり、関数の存在を知っていても楽に解けるわけではない。使うアルゴリズムがライブラリにあれば自分で実装するより楽で速くてコードが見やすい、くらいに思っていた方がよい。また、最小経路問題で距離の計算が特殊（経路中の辺の重みの最小値を距離とする等）だったりするとscipyのライブラリでは（たぶん）対応できず実装する必要あるので、ちゃんとアルゴリズムを勉強しておくのは大事。ライブラリを使えれば速いので、あまり特殊な計算ができない制約の中でどうやってライブラリを使える問題に落とし込むか、というパズル的な要素はある。が、AtCoderではscipyを使っている人がとにかく少なくて他人のコードを参考にしようにもできないことがとても多い。自分の勉強のためにも、ライブラリごり押しで解いてくれる人が増えることを願っています。