最小全域問題をプリム法で解く

最小全域木問題を学習中。プリム法を使って最小全域木を求める。応用が広いみたい。

以下の平面座標上の点の最小全域木を求める。
a = (18, 25)
b = (36, 50)
c = (36, 27)
d = (46, 37)
e = (73, 37)
f = (65, 12)

それぞれのノードが、自分を除く全てのノードへエッジをもつとする。

http://www.deqnotes.net/acmicpc/prim/

手計算で求めたらコストが101.9で辺が
a,c
c,d
d,b
d,e
e,f
だった。

# coding: utf8
#
# 平面座標上の完全グラフの最小全域木問題
#
#
from math import sqrt
INF = 10**12

class Point:
    def __init__(s, x, y):
        s.x = x
        s.y = y

def euclid_dist(u, v):
    # 2点間の距離
    return sqrt((u.x - v.x)**2 + (u.y - v.y)**2)

def MST(ps):
    # Minimum Spanning Tree

    # 距離行列の作成
    m =  [[euclid_dist(ps[i], ps[j]) for i in xrange(len(ps))] for j in xrange(len(ps))]
    # 集合の初期化
    X = set([0])
    Y = set(range(len(ps))[1:])
    #
    cost = 0
    edges = []
    print X, Y
    while len(Y) != 0:
        mcost = INF
        edge = None
        for ex in X:
            for ey in Y:
                if mcost > m[ex][ey]:
                    mcost = m[ex][ey]
                    edge = [ex, ey]
        print "[MST]: minimum path is", edge
        X.add(edge[1])
        Y.remove(edge[1])

        cost += mcost
        edges.append(edge)
    return cost, edges

def main():
    points = [[18, 25],
            [36, 50],
            [36, 27],
            [46, 37],
            [73, 37],
            [65, 12],]
    ps = [Point(*p) for p in points]
    cost, edges = MST(ps)

    vname = dict( zip(range(len(ps)), [c for c in "abcdef"]) )
    print "Cost:", cost
    print "Path:"
    for edge in edges:
        print "%s -> %s" % (vname[edge[0]], vname[edge[1]])


if __name__ == '__main__':
    main()

実行

set([0]) set([1, 2, 3, 4, 5])
[MST]: minimum cost path is [0, 2]
[MST]: minimum cost path is [2, 3]
[MST]: minimum cost path is [3, 1]
[MST]: minimum cost path is [3, 4]
[MST]: minimum cost path is [4, 5]
Cost: 101.902934864
Path:
a -> c
c -> d
d -> b
d -> e
e -> f