opengl学習 -三次元空間上で最小全域木を構築する-

前のエントリで使ったプログラムでさらに遊んでみた。というか本当はこれがやりたかった。ランダムに生成した三次元座標の集合から最小全域木を構築して表示させる。

ちなみにこれが前回やった二次元版。

そしてこれが今回作った三次元版。

実行は以下のように。

$ python mst.py 30 20 | ./a.out /dev/stdin

# coding: utf-8
#
# mst.py
#
# 3次元空間上で最小全域木を構築する
#

def distance(u, v):
    return sum(map(lambda x,y: (y-x)**2, u, v))**.5

def mst(vertices):
    edges = []
    cost  = 0

    inner = set([0])
    outer = set(xrange(1, len(vertices)))

    while len(outer) != 0:
        src = dst = None
        mindist = 1e+100
        for ii in inner: # inner index
            for oi in outer: # outer index
                dist = distance(vertices[ii], vertices[oi])
                if dist < mindist:
                    # 更新
                    src = ii
                    dst = oi
                    mindist = dist
                    #print "update:", src, dst, mindist
        cost += mindist
        inner.add(dst)
        outer.remove(dst)
        #print "path ", src,dst, mindist, cost
        edges.append([src, dst])
    return cost, edges

def main():
    #vertices = [[10., 5., -1], [10., 10., 10.], [-1., -1., -1.], [1., 1., 1.],  [-6, 1, -2],     [4., 2., 3.]]
    import random
    import sys
    if sys.argv.__len__() < 3:
        print "usage: %s npoints bound" % sys.argv[0]
        sys.exit()
    else:
        npoints = int(sys.argv[1])
        bound   = int(sys.argv[2])

    rand = lambda: random.randint(-bound, bound)
    vertices = [[rand(), rand(), rand()] for _ in xrange(npoints)]

    cost, edges = mst(vertices)
    
    for v in vertices:
        print v[0], v[1], v[2]
    print "" # new line
    for e in edges:
        print e[0], e[1]

if __name__ == '__main__':
    main()