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

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

プログラミングコンテスト攻略のためのアルゴリズムとデータ構造

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

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

隣接行列を用いるとの計算量ですが、優先度付きキューを用いることによってに計算量を抑えられます。

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))

本日のおまけ