1+ class Graph :
2+ def __init__ (self , num_of_nodes ):
3+ self .m_num_of_nodes = num_of_nodes
4+ # Initialize the adjacency matrix with zeros
5+ self .m_graph = [[0 for column in range (num_of_nodes )]
6+ for row in range (num_of_nodes )]
7+ 8+ def add_edge (self , node1 , node2 , weight ):
9+ self .m_graph [node1 ][node2 ] = weight
10+ self .m_graph [node2 ][node1 ] = weight
11+ 12+ def prims_mst (self ):
13+ # Defining a really big number, that'll always be the highest weight in comparisons
14+ postitive_inf = float ('inf' )
15+ 16+ # This is a list showing which nodes are already selected
17+ # so we don't pick the same node twice and we can actually know when stop looking
18+ selected_nodes = [False for node in range (self .m_num_of_nodes )]
19+ 20+ # Matrix of the resulting MST
21+ result = [[0 for column in range (self .m_num_of_nodes )]
22+ for row in range (self .m_num_of_nodes )]
23+ 24+ indx = 0
25+ 26+ # While there are nodes that are not included in the MST, keep looking:
27+ while (False in selected_nodes ):
28+ # We use the big number we created before as the possible minimum weight
29+ minimum = postitive_inf
30+ 31+ # The starting node
32+ start = 0
33+ 34+ # The ending node
35+ end = 0
36+ 37+ for i in range (self .m_num_of_nodes ):
38+ # If the node is part of the MST, look its relationships
39+ if selected_nodes [i ]:
40+ for j in range (self .m_num_of_nodes ):
41+ # If the analyzed node have a path to the ending node AND its not included in the MST (to avoid cycles)
42+ if (not selected_nodes [j ] and self .m_graph [i ][j ]> 0 ):
43+ # If the weight path analized is less than the minimum of the MST
44+ if self .m_graph [i ][j ] < minimum :
45+ # Defines the new minimum weight, the starting vertex and the ending vertex
46+ minimum = self .m_graph [i ][j ]
47+ start , end = i , j
48+ 49+ # Since we added the ending vertex to the MST, it's already selected:
50+ selected_nodes [end ] = True
51+ 52+ # Filling the MST Adjacency Matrix fields:
53+ result [start ][end ] = minimum
54+ 55+ if minimum == postitive_inf :
56+ result [start ][end ] = 0
57+ 58+ indx += 1
59+ 60+ result [end ][start ] = result [start ][end ]
61+ 62+ # Print the resulting MST
63+ # for node1, node2, weight in result:
64+ for i in range (len (result )):
65+ for j in range (0 + i , len (result )):
66+ if result [i ][j ] != 0 :
67+ print ("%d - %d: %d" % (i , j , result [i ][j ]))
68+ 69+ def main ():
70+ # Example graph has 9 nodes
71+ graph = Graph (9 )
72+ 73+ graph .add_edge (0 , 1 , 4 )
74+ graph .add_edge (0 , 2 , 7 )
75+ graph .add_edge (1 , 2 , 11 )
76+ graph .add_edge (1 , 3 , 9 )
77+ graph .add_edge (1 , 5 , 20 )
78+ graph .add_edge (2 , 5 , 1 )
79+ graph .add_edge (3 , 6 , 6 )
80+ graph .add_edge (3 , 4 , 2 )
81+ graph .add_edge (4 , 6 , 10 )
82+ graph .add_edge (4 , 8 , 15 )
83+ graph .add_edge (4 , 7 , 5 )
84+ graph .add_edge (4 , 5 , 1 )
85+ graph .add_edge (5 , 7 , 3 )
86+ graph .add_edge (6 , 8 , 5 )
87+ graph .add_edge (7 , 8 , 12 )
88+ 89+ graph .prims_mst ()
90+ 91+ if __name__ == "__main__" :
92+ main ()
0 commit comments