ALDS1_12_A: Minimum Spanning Tree (続き2)

解説を読む
なぜこのアルゴリズムで正しいのかは書いてないな
そこまで書いてたらボリュームがすごいことになるってことだろう

このアルゴリズムを実装するポイントは、辺の選択ステップで"どのように最小の重みをもつ辺を保存しておくか"です。

そこだなきっと

puのpはparentのpだったか
MSTもTreeだもんな
自分のコードは辺をバラバラに覚えてるだけで、木の形にはなってなかった
情報としては同じだけ持ってると思うけど

しかしやっぱりWHITE、GRAY、BLACKか
だんだんやることが複雑になってくると必要になってくるのかも？

テキストのC++コードを参考にして書き直し
ほぼ全面書き換え

ん、配列が0オリジンだな
どこでつじつまを合わせてるんだと思ったけどどこでも合わせてない
そもそもi,jが1からとは書いてないから自分で0オリジンと思って書けばいいだけか

color[i]がBLACKになっているということが
頂点iがTに含まれるっていうことを意味してるわけだな

from sys import maxsize

WHITE = 0
GRAY = 1
BLACK = 2

def read_adj_mat():
    n = int(input())
    M = []
    for _ in range(n):
        M.append([maxsize if x == "-1" else int(x)
                  for x
                  in input().split()])
    return n, M

def min_sp_tree(n, M):
    d = [maxsize] * n
    p = [-1] * n
    color = [WHITE] * n

    d[0] = 0

    while True:
        minv = maxsize
        u = -1

        for i in range(n):
            if minv > d[i] and color[i] != BLACK:
                u = i
                minv = d[i]

        if u == -1:
            break

        color[u] = BLACK
        for v in range(n):
            if color[v] != BLACK and M[u][v] != maxsize:
                if d[v] > M[u][v]:
                    d[v] = M[u][v]
                    p[v] = u
                    color[v] = GRAY

    return sum([M[i][p[i]] for i in range(n) if p[i] != -1])

def main():
    n, M = read_adj_mat()
    print(min_sp_tree(n, M))

if __name__ == '__main__':
    main()

100x100で00.03s
やっぱこっちのほうが速い