2024.06.23追記: 2023年の言語アップデート後、networkxの実行時間が大きく改善されたことがわかったため、記事を大幅に更新しました。
2020年のAtCoder言語アップデートにより、Python 3.8環境ではnetworkxが使えるようになりました。
networkxを使うことで自前実装に比べてコードを非常に簡潔に書くことができ、強力な武器になりえます。
PAST-O これで通るのか。 networkx ずるくない？ pic.twitter.com/Dvvq8a6EwG
— きり (@kiri8128) June 6, 2020
一方でnetworkxはメモリ使用量・実行速度ともにふんだんに使っていく傾向があり、自前実装に比べても比較的TLE/MLEしがちであることが分かっています。
この記事ではnetworkxの基本操作についてまとめた後、どのような問題なら現実的に利用可能か確認します。
- 本文にはAtCoder過去問の解説/実装が含まれますのでご理解ください。
基本的な操作
宣言
import networkx as nx
G = nx.Graph() # 単純無向グラフ
DG = nx.DiGraph() # 単純有向グラフ
頂点の追加
G.add_node(1) # 頂点を一つ追加
G.add_nodes_from((2, 3, 4)) # 頂点を一括で追加
G.add_nodes_from(range(1, N + 1)) # 1からNまでの頂点を一括で追加
辺の追加
G.add_edge(1, 2) # 辺を一つ追加。頂点が足りなければ自動で追加される
G.add_edges_from([(1, 3), (2, 4)]) # 一括で追加
有向グラフの場合、$from,to$の順に入れます。
add_edges_fromを使うと、例えば以下のような入力
\begin{align}
&N \\
&x_1 \ \ y_1 \\
&x_2 \ \ y_2 \\
&\vdots \\
&x_N\ \ y_N
\end{align}
に対して、
N = int(input())
G = nx.Graph()
G.add_edges_from([tuple(map(int, input().split())) for _ in range(N)])
のように書くことができます。
2020.9.6 追記：下記のように書くこともできますが、速度的に良くはならないようです。(後述)
import networkx as nx
N = 2 * 10 ** 5
G = nx.Graph([tuple(map(int, input().split())) for _ in range(N)])
重み付き辺の追加
G.add_edge(1, 2, weight=3) # 一つ追加するだけだとキーワード引数で指定しないとダメ
G.add_weighted_edges_from([(1, 3, 10),(2, 4, 15)] # 一括で追加
add_weighted_edges_fromでは、$from,to,weight$の順に入れます。AtCoderでの標準入力は多くの場合この順番なので、そのまま[tuple(map(int, input().split())) for _ in range(N)]で入れることができますね。
で、実はadd_weighted_edges_fromに限ってはtupleを外して書ける([map(int, input().split()) for _ in range(N)]とできる)ようです。これが仕様なのかちょっとよく分かっていないので情報お待ちしています。
辺・頂点の列挙
for n in G.nodes(): # 頂点n
print(n)
for x, y in G.edges(): # 辺(x,y)
print(x, y)
for x, y, d in G.edges(data=True): # 重み付き辺がいる場合。dは辞書なので注意
weight = d['weight']
print(x, y, weight)
基本的な操作の速度検証
AtCoderコードテスト環境での検証です。
インポート
import networkx as nx
標準入力：なし
| ① | ② | ③ | ④ | ⑤ | ⑥ | ⑦ | ⑧ | ⑨ | ⑩ | |
|---|---|---|---|---|---|---|---|---|---|---|
| 実行速度 | $\mathrm{104\ ms}$ | $\mathrm{110\ ms}$ | $\mathrm{106\ ms}$ | $\mathrm{105\ ms}$ | $\mathrm{103\ ms}$ | $\mathrm{104\ ms}$ | $\mathrm{110\ ms}$ | $\mathrm{105\ ms}$ | $\mathrm{106\ ms}$ | $\mathrm{104\ ms}$ |
参考までに、CPython 3.11.4 での import numpy as np は 約 $\mathrm{80\ ms}$程度です。
インポート+宣言
import networkx as nx
G = nx.Graph()
標準入力：なし
| ① | ② | ③ | ④ | ⑤ | ⑥ | ⑦ | ⑧ | ⑨ | ⑩ | |
|---|---|---|---|---|---|---|---|---|---|---|
| 実行速度 | $\mathrm{104\ ms}$ | $\mathrm{107\ ms}$ | $\mathrm{106\ ms}$ | $\mathrm{106\ ms}$ | $\mathrm{108\ ms}$ | $\mathrm{108\ ms}$ | $\mathrm{110\ ms}$ | $\mathrm{105\ ms}$ | $\mathrm{105\ ms}$ | $\mathrm{105\ ms}$ |
Digraphでもやりましたがほとんど同じ感じでした。誤差といっていいと思います。
インポート+宣言+頂点追加
頂点数$N=2 \times 10^5$で試します。ダイクストラ法などで一般的な制約だと思います。
import networkx as nx
N = 2 * 10 ** 5
G = nx.Graph()
G.add_nodes_from(range(1, N + 1))
標準入力：なし
| ① | ② | ③ | ④ | ⑤ | ⑥ | ⑦ | ⑧ | ⑨ | ⑩ | |
|---|---|---|---|---|---|---|---|---|---|---|
| 実行速度 | $\mathrm{221\ ms}$ | $\mathrm{220\ ms}$ | $\mathrm{220\ ms}$ | $\mathrm{226\ ms}$ | $\mathrm{232\ ms}$ | $\mathrm{241\ ms}$ | $\mathrm{219\ ms}$ | $\mathrm{220\ ms}$ | $\mathrm{217\ ms}$ | $\mathrm{216\ ms}$ |
頂点の追加は$10^5$個あたり$\mathrm{60\ ms}$程度とみてよさそうです。
インポート+宣言+頂点追加+辺追加
頂点数$N=2 \times 10^5$、辺数$M=2 \times 10^5 - 1$の直線グラフで試します(もうちょい複雑なグラフの方がいいかもしれませんが生成が面倒なので)。
import networkx as nx
N = 2 * 10 ** 5
E = [(i, i + 1) for i in range(1, N)]
G = nx.Graph()
G.add_nodes_from(range(1, N + 1))
G.add_edges_from(E)
標準入力：なし
| ① | ② | ③ | ④ | ⑤ | ⑥ | ⑦ | ⑧ | ⑨ | ⑩ | |
|---|---|---|---|---|---|---|---|---|---|---|
| 実行速度 | $\mathrm{498\ ms}$ | $\mathrm{489\ ms}$ | $\mathrm{477\ ms}$ | $\mathrm{491\ ms}$ | $\mathrm{491\ ms}$ | $\mathrm{497\ ms}$ | $\mathrm{475\ ms}$ | $\mathrm{485\ ms}$ | $\mathrm{496\ ms}$ | $\mathrm{505\ ms}$ |
辺の追加は$10^5$本あたり$\mathrm{60\ ms}$程度とみてよさそうです。実行時間 $2 \sec$の問題ならすでに1/4を使い切ってしまいました。
※ところで、networkxにはグラフに辺をつけるとき頂点がなければ自動で追加する機能があります。これを利用して頂点の宣言を省略すると、コード量は減りますが、実行時間はわずかに伸びます。
import networkx as nx
N = 2 * 10 ** 5
E = [(i, i + 1) for i in range(1, N)]
G = nx.Graph()
G.add_edges_from(E)
標準入力：なし
| ① | ② | ③ | ④ | ⑤ | ⑥ | ⑦ | ⑧ | ⑨ | ⑩ | |
|---|---|---|---|---|---|---|---|---|---|---|
| 実行速度 | $\mathrm{497\ ms}$ | $\mathrm{524\ ms}$ | $\mathrm{538\ ms}$ | $\mathrm{532\ ms}$ | $\mathrm{559\ ms}$ | $\mathrm{499\ ms}$ | $\mathrm{591\ ms}$ | $\mathrm{552\ ms}$ | $\mathrm{542\ ms}$ | $\mathrm{595\ ms}$ |
さらに、add_edges_fromを使わない場合は大幅に速度が低下します。
import networkx as nx
N = 2 * 10 ** 5
G = nx.Graph([(i, i + 1) for i in range(1, N)])
標準入力：なし
| ① | ② | ③ | ④ | ⑤ | ⑥ | ⑦ | ⑧ | ⑨ | ⑩ | |
|---|---|---|---|---|---|---|---|---|---|---|
| 実行速度 | $\mathrm{727\ ms}$ | $\mathrm{772\ ms}$ | $\mathrm{753\ ms}$ | $\mathrm{748\ ms}$ | $\mathrm{758\ ms}$ | $\mathrm{845\ ms}$ | $\mathrm{739\ ms}$ | $\mathrm{807\ ms}$ | $\mathrm{829\ ms}$ | $\mathrm{799\ ms}$ |
インポート+宣言+頂点追加+重み付き辺追加
頂点数$N=2 \times 10^5$、辺数$M=2 \times 10^5 - 1$、重みはすべて$2$の直線グラフで試します。
import networkx as nx
N = 2 * 10 ** 5
E = [(i, i + 1, 2) for i in range(1, N)]
G = nx.Graph()
G.add_nodes_from(range(1, N + 1))
G.add_weighted_edges_from(E)
標準入力：なし
| ① | ② | ③ | ④ | ⑤ | ⑥ | ⑦ | ⑧ | ⑨ | ⑩ | |
|---|---|---|---|---|---|---|---|---|---|---|
| 実行速度 | $\mathrm{519\ ms}$ | $\mathrm{541\ ms}$ | $\mathrm{511\ ms}$ | $\mathrm{523\ ms}$ | $\mathrm{536\ ms}$ | $\mathrm{510\ ms}$ | $\mathrm{504\ ms}$ | $\mathrm{528\ ms}$ | $\mathrm{523\ ms}$ | $\mathrm{532\ ms}$ |
重み付き辺の追加は$10^5$本あたり$\mathrm{70\ ms}$程度でした。
このように、networkxのグラフは生成が重いため、頂点数の巨大な問題はそれだけでかなり困難さが増しそうです。頂点数・辺数が小さい問題での利用が確実性が高いと思います。
単純グラフ
ダイクストラ法
ABC035 D - トレジャーハント
問題概要
$N$ 頂点 $M$ 辺の重み付き有向グラフが与えられます。ここでグラフの重みは、その辺を移動するのにかかる所要時間です。
頂点$i$に居ると毎分$A_i$円もらえます。時刻$0$に頂点$1$を出発し、時刻$T$には戻ってきている時所持金を最大化してください。
| 頂点数 $N$ | 辺数 $M$ | TL | ML | 提出例 | 提出解の時間計算量 | 実行結果 |
|---|---|---|---|---|---|---|
| $10^5$ | $\min(N(N-1)/2,10^5)$ | $3 \sec$ | $256\ \mathrm{MB}$ | networkx① | $O((N+M)\log N)$ | $\mathrm{1862\ ms}$ |
| networkx② | 同上 | $\mathrm{1085\ ms}$ | ||||
| scipy | 同上 | $\mathrm{311\ ms}$ |
解法概要
滞在する頂点は1つでよいので、各頂点$i$に対して$1 \to i \to 1$と移動したときの頂点$i$に居られる時間を求めることで、どの頂点を目指すのが最適か分かります。
頂点$i$に居られる時間は、$T-(往路1\to iの所要時間)-(復路i \to 1の所要時間)$です。往路の所要時間は頂点$1$始点のダイクストラ法で、復路の所要時間は有向辺の向きを逆にしたグラフにおける頂点$1$始点のダイクストラ法で求められるので、どちらも$O((N+M)\log N)$が実現できます。
実装
有向グラフなのでnx.DiGraph()を使います。有向辺の向きを逆にしたグラフの生成には、DiGraph.reverseというおあつらえ向きの関数があったのでそれを使います。
次にダイクストラ法パートです。一口にダイクストラ法といっても色々な関数があるので是非公式リファレンスを見ていただきたいのですが、今回は「起点1つ、終点は全頂点、経路の構築不要」なのでsingle_source_dijkstra_path_lengthを使いました。
import networkx as nx
N, M, T = map(int, input().split())
A = tuple(map(int, input().split()))
G = nx.DiGraph()
G.add_nodes_from(range(1, N + 1))
G.add_weighted_edges_from([map(int, input().split()) for _ in range(M)])
G_inv = G.reverse() # 有向辺の向きを逆にしたグラフを生成
L = nx.single_source_dijkstra_path_length(G, 1) # {頂点名: 距離} の辞書型で返される
L_inv = nx.single_source_dijkstra_path_length(G_inv, 1)
ans = 0
for i in range(1, N + 1):
# 頂点1と目標の頂点とが連結の場合のみkeyがある
if i in L and i in L_inv:
ans = max(ans, (T - L[i] - L_inv[i]) * A[i - 1])
print(ans)
上記のコードでもACが得られますが、メモリ節約の工夫をします。先ほどDiGraph.reverseという関数を使って有向辺の向きを逆転させたグラフをつくりましたが、この関数のキーワード引数をcopy=Falseとすることでグラフのビューを作るようになり再生成を避けられます。
# 前略
G = nx.DiGraph()
G.add_nodes_from(range(1, N + 1))
G.add_weighted_edges_from([map(int, input().split()) for _ in range(M)])
L = nx.single_source_dijkstra_path_length(G, 1)
G_inv = G.reverse(copy=False) # ここを変更。ついでに往路のdijkstra法実行は前に持っていった
L_inv = nx.single_source_dijkstra_path_length(G_inv, 1)
# 後略
これで、メモリ使用量だけでなくなんと実行時間も$1/2$近くなりました。
実行時間
頂点・辺が付与されたグラフをなるべく減らす・サイズを小さくすることがミソのようです。
しかしダイクストラ法に関してはscipyにもありますし、場合によっては自前実装したほうが高速で確実というケースもあるようです。一般的なグラフの問題でわざわざnetworkxのダイクストラ法メゾットを使う必要があるかは、微妙と言わざるを得ないと思います。
フロイドワーシャル法
ABC051 D - Candidates of No Shortest Paths
問題概要
$N$ 頂点 $M$ 辺の重み付き単純無向連結グラフが与えられます。どの異なる 2 頂点間の、どの最短経路にも含まれない辺の数を求めてください。
| 頂点数 $N$ | 辺数 $M$ | TL | ML | 提出例 | 提出解の時間計算量 | 実行結果 |
|---|---|---|---|---|---|---|
| $100$ | $\min(N(N-1)/2,1000)$ | $2 \sec$ | $256\ \mathrm{MB}$ | networkx | $O(N^3+MN^2)$ | $\mathrm{1259\ ms}$ |
解法概要
まず全頂点間の最短経路長をフロイドワーシャル法で求めます。次に、各辺$(i,j)$に対して、
dist(s,i)+edge(i,j)+dist(j,t)=dist(s,j)
を満たす$(s,t)$がないか全探索します(ただし$dist$は二点間の距離、$edge$は辺の長さを表す)。もし存在すれば、辺$(i,j)$はこの$(s,t)$間の最短経路に含まれます。
実装
無向グラフなのでnx.Graph()を使います。
フロイドワーシャル法の関数もいくつかありますが(公式リファレンス)、ここは普通のfloyd_warshall関数を使います。
次に辺と頂点を全探索するわけですが、列挙も下にある通りそれっぽい感じで上手くいきます。
(networkxとは関係ないですが、itertools.combinationsはこの場合$N$頂点のうちから$2$頂点を順不同で選ぶ組み合わせを全列挙します)
import networkx as nx
import itertools
N, M = map(int, input().split())
G = nx.Graph()
G.add_nodes_from(range(1, N + 1))
G.add_weighted_edges_from([map(int, input().split()) for _ in range(M)])
S = nx.floyd_warshall(G)
ans = M
for i, j, d in G.edges(data=True):
for s, t in itertools.combinations(G.nodes(), 2):
if S[s][i] + d['weight'] + S[j][t] == S[s][t]:
ans -= 1
break
print(ans)
実行時間
こちらも$10^6$くらいの計算回数がありますが、1sほど使ってACできています。
あと本問に関しては、editorialによると全探索パートを高速化することができるようです。
その他
ABC075 C - Bridge
問題概要
$N$ 頂点 $M$ 辺の単純無向連結グラフが与えられます。与えられた $M$本のうち橋である辺の本数を求めてください。
| 頂点数 $N$ | 辺数 $M$ | TL | ML | 提出例 | 提出解の時間計算量 | 実行結果 |
|---|---|---|---|---|---|---|
| $50$ | $\min(N(N-1)/2,50)$ | $2 \sec$ | $256\ \mathrm{MB}$ | networkx | $O(N+M)$1 | $\mathrm{117\ ms}$ |
実装
橋を列挙する関数があったので(！)、それを貼ります。
import networkx as nx
N, M = map(int, input().split())
G = nx.Graph()
G.add_nodes_from(range(1, N + 1))
G.add_edges_from([tuple(map(int, input().split())) for _ in range(M)])
print(len(tuple(nx.bridges(G))))# nx.bridges……橋となる辺をイテレータとして返す
実行時間
想定解が$O(M(N+M))$だったようで、充分な時間で間に合っています(本当にnx.bridgeの時間計算量が$O(N+M)$なのか怪しい気もしますが……)。
さすがに今後こういう貼るだけが出ることはないでしょう。
グリッドグラフ
グリッドグラフを生成する関数としてnx.grid_graphが、特に二次元グリッドグラフを生成する関数としてnx.grid_2d_graphがあります。
nx.grid_2d_graphを使ってグラフを生成した場合、グラフの頂点名は自然に$(h,w)$のtupleになります。頂点名をこのようなtupleにできるのはとても扱いやすくて良さそうです。
ABC020 C - 壁抜け
問題概要
$H \times W$のグリッド状に並んでいるマス上で、スタートからゴールまで行きたいです。マスには白マス(スタート・ゴールを含む)と黒マスがあり、白マスに入るには$1$秒、黒マスに入るには$x$秒かかります($x$はこちらが指定できます)。$T$秒以内にゴールにたどり着けるような$x$の最大値を求めてください。
| 頂点数 $N(=H \times W)$ | 辺数 $M(=2H(W-1)+2(H-1)W)$ | TL | ML | 提出例 | 提出解の時間計算量 | 実行時間 |
|---|---|---|---|---|---|---|
| $100$ | $180$ | $2 \sec$ | $256\ \mathrm{MB}$ | networkx | $O((N+M)\log N \log T)$ | $\mathrm{137\ ms}$ |
| scipy | 同上 | $\mathrm{134\ ms}$ |
解法概要
二分探索により$x$の最大値を探ります。$x$が固定されていれば、何秒でたどり着けるかはダイクストラ法により判定できます。
実装
nx.grid_graphで生成したグラフの辺には重みが登録されていません。重みが登録されていない辺は自動的に重み$1$としてダイクストラ法メゾットが実行されることを利用し、重み$x$の辺だけに重みを登録していきます。
import networkx as nx
H, W, T = map(int, input().split())
S = [input() for _ in range(H)]
for h in range(H):
for w in range(W):
if S[h][w] == 'S':
start = (h, w)
elif S[h][w] == 'G':
goal = (h, w)
def canArrive(x):
'''判定関数
'''
G = nx.DiGraph(nx.grid_2d_graph(H, W))
for h in range(H):
for w in range(W):
# 黒に入る辺の重みをxに変更
if S[h][w] == '#':
for dh, dw in ((1, 0), (0, 1), (-1, 0), (0, -1)):
if 0 <= h - dh < H and 0 <= w - dw < W:
G[(h - dh, w - dw)][(h, w)]['weight'] = x
return nx.dijkstra_path_length(G, start, goal) <= T
# 二分探索
l, r = 1, 10 ** 9
while r - l > 1:
m = (l + r) // 2
if canArrive(m):
l = m
else:
r = m
print(l)
実行時間
こちらは$10^3$～$10^4$くらいの計算回数で、そもそもかなり実行時間がゆるいといえます。また CPython 3.11.4 ではscipyでの実装とほぼ同じ実行時間を達成しているのも注目すべきでしょう。
ABC151 D - Maze Master
問題概要
$H \times W$のグリッド状に並んでいるマス上で、”道”のマスだけを通ってスタートからゴールまで最短経路で行きます。このような最短経路が最も長くなるようなスタートからゴールを探し、その道のりの長さを求めてください。
| 頂点数 $N(=H \times W)$ | 辺数 $M(=2H(W-1)+2(H-1)W)$ | TL | ML | 提出例 | 提出解の時間計算量 | 実行時間 |
|---|---|---|---|---|---|---|
| $400$ | $1520$ | $2 \sec$ | $1024\ \mathrm{MB}$ | networkx① | $O(N^2)$ | $\mathrm{296\ ms}$ |
| networkx② | 同上 | $\mathrm{293\ ms}$ | ||||
| 素Python | 同上 | $\mathrm{118\ ms}$ |
解法概要
各始点についてbfsをしてそれぞれ最も遠い点を探します。
実装
グリッドグラフから”壁”に入る辺を削除する実装にしてみました。
BFS用の関数nx.bfs_edges()があったので使います。使いどころは限られるかもしれませんが、そこそこ便利です。
import networkx as nx
H, W = map(int, input().split())
maze = [input() for _ in range(H)]
G = nx.grid_2d_graph(H, W)
for h in range(H):
for w in range(W):
# 壁に入る辺を削除
if maze[h][w] == '#':
for dh, dw in ((1, 0), (0, 1), (-1, 0), (0, -1)):
if ((h - dh, w - dw), (h, w)) in G.edges():
G.remove_edge((h - dh, w - dw), (h, w))
def bfs(sy, sx):
d = dict()
d[(sy, sx)] = 0
for coordinate_from, coordinate_to in nx.bfs_edges(G, (sy, sx)):
d[coordinate_to] = d[coordinate_from] + 1
return max(d.values())
ans = 0
for y in range(H):
for x in range(W):
if maze[y][x] == '.':
ans = max(ans, bfs(y, x))
print(ans)
実行時間
こちらは$10^5$くらいの計算回数です。グラフも小さいですし、十分通りますね。
上にコードを載せたのは提出例のnetworkx①です。nx.grid_2d_graphを紹介したかったのでこういう実装になりましたが、普通に"道"どうしをつなぐ辺を加えていく実装でもよいです( networkx② )。②の方が実行速度は速くなるかな？と思っていましたがほぼ同じでした。
木
ABC138 D - Ki
問題概要
$N$頂点の木が与えられます。根は頂点$1$とします。各頂点にカウンターがあり、はじめはカウンターの値は全ての頂点で$0$です。
以下の$Q$回のクエリを実行します。
- ある頂点$p_j$を根とする部分木に含まれるすべての頂点にあるカウンターに$x_j$を足す。
全てのクエリが終わった時、各頂点のカウンターの値を求めてください。
| 頂点数 $N$ | TL | ML | 提出例 | 提出解の時間計算量 | 実行時間 |
|---|---|---|---|---|---|
| $2 \times 10^5$ | $2 \sec$ | $1024\ \mathrm{MB}$ | networkx | $O(Q+N)$ | TLE |
| 素Python | 同上 | $\mathrm{875\ ms}$ |
解法概要
累積和です。
実装
頂点をたどる順序は多分なんでもいいんですが、今回はDFS順にします。DFS用の関数nx.dfs_edgesを使ってイテレートし、累積和をとっていきます。
import networkx as nx
int1 = lambda x: int(x) - 1 # 0-indexedに直す関数
N, Q = map(int, input().split())
G = nx.Graph()
G.add_nodes_from(range(N))
G.add_edges_from([tuple(map(int1, input().split())) for _ in range(N - 1)])
lst = [0] * N
for _ in range(Q):
p, x = map(int, input().split())
lst[p - 1] += x
# 累積和
for parent, child in nx.dfs_edges(G, 0):
lst[child] += lst[parent]
print(*lst)
実行時間
頂点数が$2 \times 10^5$と巨大なので、やはり厳しかったようです。Pythonでも$0.8 \sec$以上かかっています。
ABC133 E - Virus Tree
問題概要
$N$頂点の木が与えられます。$K$色までの色を使って下記を満たすように木の頂点を塗り分けるとき、塗分け方を数え上げて$10^9+7$で割った余りを答えてください。
- 二つの異なる頂点 $x,y$間の距離が $2$ 以下ならば、頂点 $x$ の色と頂点 $y$ の色は異なる。
| 頂点数 $N$ | TL | ML | 提出例 | 提出解の時間計算量 | 実行時間 |
|---|---|---|---|---|---|
| $10^5$ | $2 \sec$ | $1024\ \mathrm{MB}$ | networkx | $O(N)$ | $\mathrm{1333\ ms}$ |
| 素Python | 同上 | $\mathrm{278\ ms}$ |
解法概要
適当に根を決めて根付き木とみなしたとき、各頂点$x$に対して以下が成り立ちます。
- 頂点$x$の子の数を$c_x$と置くと、子の塗分け方は${}_{K-2}P_{c_x}$通り。
あとは根に親がないことに注意して実装します。
実装
これもDFS順でたどってみます。nx.degree()で各頂点の頂点名と次数とが対応した辞書が手に入ります。
import networkx as nx
int1 = lambda x: int(x) - 1
MOD = 10 ** 9 + 7
N, K = map(int, input().split())
G = nx.Graph()
G.add_nodes_from(range(N))
G.add_edges_from([tuple(map(int1, input().split())) for _ in range(N - 1)])
deg = nx.degree(G)
ans = K
for i in range(deg[0]):
ans = (ans * (K - i - 1)) % MOD
for parent, child in nx.dfs_edges(G, 0):
for i in range(deg[child] - 1):
ans = (ans * (K - i - 2)) % MOD
print(ans)
実行速度
これもヤバいかなと思っていましたが、案外TLEしませんでした。
フロー
最小カット
競プロ典型90問 040 - Get More Money
問題概要
$N$軒の家があり、各家に入るか入らないかを決めてもらえるお金を最大化したいです。
- 家 $i$ に入ると $A_i - W$ 円手に入ります。
- 家 $i$ に入ると ほかの家への鍵 $c_{i,1} \ldots c_{i,k_i}$ が手に入ります。
- 家 $j$ に入るには、先にほかの家にある鍵 $j$ をすべて集めないといけません。
| 頂点数 $N$ | TL | ML | 提出例 | 提出解の時間計算量 | 実行時間 |
|---|---|---|---|---|---|
| $100$ | $2 \sec$ | $1024\ \mathrm{MB}$ | networkx | $O(N^4)$ | $\mathrm{136\ ms}$ |
解法概要
Project Selection Problem （いわゆる「燃やす埋める問題」）に帰着され、うまく有向グラフを構成して最小カットを求める問題となります。
実装
networkxのフローライブラリはかなり充実しており、今回はnx,minimum_cut_value()を用います。フローでは辺容量制約を考えることが多いですが、capacityによってこれを表現することが出来ます。
import networkx as nx
N, W = map(int, input().split())
A = list(map(int, input().split()))
G = nx.DiGraph()
G.add_node("S")
G.add_node("T")
G.add_nodes_from(range(1, N + 1))
INF = 10**18
for i, a in enumerate(A, 1):
G.add_edge(i, "T", capacity=a)
G.add_edge("S", i, capacity=W)
for i in range(1, N + 1):
k, *c = map(int, input().split())
for x in c:
G.add_edge(i, x, capacity=INF)
print(sum(A) - nx.minimum_cut_value(G, "S", "T"))
実行速度
余裕でした。
フローはグラフサイズが小さい問題が多いため、networkxとの相性が良い印象があります。
おわりに
結論としては、
- 自前実装より早くなることはまずない
- $O(N),O(N \log N)$は頂点数$10^5$、$O(N^3)$は頂点数$100$くらいは通る、ただし前後の計算が重いとTLEする場合もある。
といったところでしょうか。コンテストで使うとしたら、実行時間にゆとりがありそうだと見切った時に、実装を簡潔にするために使うくらいの付き合い方になるかなと感じました。高速化テクが発明されることを願っています。