徒然weed

アウトプットの場

ゆる〜くpythonの練習「最小全域木編」

f:id:shintaro-football7:20191022195746p:plain

今回は最小全域木についてです。

問題はこちら。
http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ALDS1_12_A

「プリムのアルゴリズム」を用います。
プリムのアルゴリズムとは既にカバーしている点の集合をT、それ以外をV-TとしてTからV-Tに伸びる最小の辺を見つけてTをおきくしていくアルゴリズムです。
f:id:shintaro-football7:20191030100452p:plain

隣接行列を用いるとO({|V|^2})の計算量ですが、優先度付きキューを用いることによってO({Vlog V})に計算量を抑えられます。

from heapq import heappush, heappop, heapify

n = int(input())
G = [[] for _ in range(n)]

for i in range(n):
    for j,dist in enumerate(map(int,input().split())):
        if dist != -1:
            G[i].append((j,dist))
sum = 0

def prim(n):
    que = [(w,c) for c,w in G[0]]
    heapify(que)
    color = [0]*n
    color[0] = 1
    global sum
    while que:
        w,node = heappop(que)
        if color[node] == 1:
            continue
        color[node] = 1
        sum += w
        for x,y in G[node]:
            if color[x] == 1:
                continue
            heappush(que,(y,x))
    return sum

print(prim(n))

本日のおまけ