-
Notifications
You must be signed in to change notification settings - Fork 4
/
mst_kruskal.py
57 lines (35 loc) · 1.37 KB
/
mst_kruskal.py
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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
import sys
sys.path.append("..")
"""
In this programming problem you'll code up Prim's minimum spanning tree algorithm.
This file describes an undirected graph with integer edge costs. It has the format
[number_of_nodes] [number_of_edges]
[one_node_of_edge_1] [other_node_of_edge_1] [edge_1_cost]
[one_node_of_edge_2] [other_node_of_edge_2] [edge_2_cost]
You should NOT assume that edge costs are positive, nor should you assume that they are distinct.
Your task is to run Prim's minimum spanning tree algorithm on this graph.
You should report the overall cost of a minimum spanning tree --- an integer, which may or may not be negative
"""
from utils.file_reader import *
from union_find import UnionFind
def get_data():
lines = list(map(lambda s: s.split(), read_input(FilePath.MST_PRIM)))
num_nodes, num_edges = map(int, lines[0])
g = []
for line in lines[1:]:
u, v, cost = map(int, line)
g.append((cost, u, v))
return g, num_nodes
def kruskal_mst(g, num_nodes):
total_cost = index = 0
g.sort()
uf = UnionFind(num_nodes)
while uf.size() > 1:
cost, u, v = g[index]
if uf.union(u, v):
total_cost+=cost
index+=1
return total_cost
g, num_nodes = get_data()
prims_mst_total_cost = kruskal_mst(g, num_nodes)
print('Kruskal MST total cost:', prims_mst_total_cost)