Algorithms/prims.py at master · Manikumar1998/Algorithms · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
import sys

def min_node(graph, keys, mst):
    minimum = sys.maxint
    min_index = None
    for i in xrange(len(keys)):
        if mst[i] == False and keys[i] < minimum:
            minimum = keys[i]
            min_index = i
    return min_index

def prims(graph, N):
    mst = [False for i in xrange(N)]
    keys = [sys.maxint for i in xrange(N)]
    edges = []

    keys[0] = 0

    new_index = min_node(graph, keys, mst)
    while(new_index != None):
        mst[new_index] = True

        for i in xrange(len(graph)):
            if graph[new_index][i] < keys[i] and graph[new_index][i] != 0:
                keys[i] = graph[new_index][i]
                edges.append([new_index, i, keys[i]])
        new_index = min_node(graph, keys, mst)
    print edges

if __name__ == '__main__':
    graph = [[0, 2, 0, 6, 0],
             [2, 0, 3, 8, 5],
             [0, 3, 0, 0, 7],
             [6, 8, 0, 0, 9],
             [0, 5, 7, 9, 0]]
    prims(graph, len(graph))