Prim.py

WHITE = 0
GRAY  = 1
BLACK = 2
INF   = 1<<50

def prim(n):
    """
    n * n 行列の隣接行列に対して
    """
    d = [INF]*n
    p = [-1]*n
    color = [WHITE]*n
    d[0] = 0
    while True:
        minv = INF
        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] != INF:
                if d[v] > M[u][v]:
                    d[v] = M[u][v]
                    p[v] = u
                    color[v] = GRAY
    sum = 0
    for i in range(n):
        if p[i] != -1:
            sum += M[i][p[i]]
    return sum #最小全域木の大きさ


# input
M = []
for line in open("p107.txt","r"):
    M.append([int(i) for i in line.replace("-","0").split(",")])

#グラフの辺の合計を求める
a = 0
for i in range(len(M)):
    for j in range(i+1): #iまで 
        a += M[i][j]

#0 を INF に変換
for i in range(len(M)):
    for j in range(len(M[i])):
        if M[i][j] == 0:
            M[i][j] = INF


print("最小全域木の大きさは",prim(len(M)))
print("グラフの辺の大きさの合計は",a)
print("ans",a-prim(len(M)))