Pythonで理解する蟻本「2-5 最小全域木問題１（プリム法）」(p.100)
この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-5 最小全域木問題１（プリム法）」(p.100)
のコードをPythonで書き直したものとなっています。
入力
入力例
7 9
0 1 2
0 4 10
1 2 1
1 3 3
1 5 7
3 5 1
3 6 5
4 5 5
5 6 8
蟻本に載っているコード（O(|V|^2)）
INF = float('inf')
V, E = map(int,input().split()) # Vは頂点数、Eは辺数
cost = [[INF] * V for _ in range(V)] # cost[u][v]は辺e=(u,v)のコスト(存在しない場所はINF)
mincost = [INF] * V # 集合Xからの辺の最終コスト
used = [False] * V # 頂点iがXに含まれているか
for _ in range(E):
s, t, c = map(int,input().split())
cost[s][t] = c
cost[t][s] = c
def prim():
mincost[0] = 0
res = 0
while True:
v = -1
# Xに属さない頂点のうちXからの辺のコストが最小になる頂点を探す
for u in range(V):
if not used[u] and (v == -1 or mincost[u] < mincost[v]):
v = u
if v == -1:
break
used[v] = True # 頂点vをXに追加
res += mincost[v] # 辺のコストを加える
for u in range(V):
mincost[u] = min(mincost[u], cost[v][u])
return res
res = prim()
print(res)
実行結果
17
優先度付きキューを使った解法（O(|E| log |V|)）
# 優先度付きキューをインポートする
import heapq
INF = float('inf')
V, E = map(int,input().split()) # Vは頂点数、Eは辺数
cost = [[INF] * V for _ in range(V)] # cost[u][v]は辺e=(u,v)のコスト(存在しない場所はINF)
used = [False] * V # 頂点iがXに含まれているか
for _ in range(E):
s, t, c = map(int,input().split())
cost[s][t] = c
cost[t][s] = c
def prim():
res = 0
q = [[0, 0]]
while True:
v = -1
MIN = INF
# Xに属さない頂点のうちXからの辺のコストが最小になる頂点を探す
while len(q) != 0:
c, u = heapq.heappop(q)
if not used[u] and (v == -1 or c < MIN):
v = u
MIN = c
break
if v == -1:
break
used[v] = True # 頂点vをXに追加
res += c # 辺のコストを加える
for u in range(V):
if cost[v][u] != INF:
heapq.heappush(q, [cost[v][u], u])
return res
res = prim()
print(res)
実行結果
17
