 # MST_Prim

Aug 9th, 2022
545
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. from queue import PriorityQueue
2.
3.
4. class Edge:
5.     def __init__(self, first, second, weight):
6.         self.first = first
7.         self.second = second
8.         self.weight = weight
9.
10.     def __gt__(self, other):
11.         return self.weight > other.weight
12.
13.
14. def prim(node, graph, forest, forest_edges):
16.     pq = PriorityQueue()
17.     for edge in graph[node]:
18.         pq.put(edge)
19.
20.     while not pq.empty():
21.         min_edge = pq.get()
22.         none_tree_node = -1
23.
24.         if min_edge.first in forest and min_edge.second not in forest:
25.             none_tree_node = min_edge.second
26.         elif min_edge.first not in forest and min_edge.second in forest:
27.             none_tree_node = min_edge.first
28.
29.         if none_tree_node == -1:
30.             continue
32.         forest_edges.append(min_edge)
33.
34.         for edge in graph[none_tree_node]:
35.             pq.put(edge)
36.
37.
38. edges = int(input())
39.
40. graph = {}
41.
42. for _ in range(edges):
43.     first, second, weight = [int(x) for x in input().split(', ')]
44.     if first not in graph:
45.         graph[first] = []
46.     if second not in graph:
47.         graph[second] = []
48.     edge = Edge(first, second, weight)
49.     graph[first].append(edge)
50.     graph[second].append(edge)
51.
52. forest = set()
53. forest_edges = []
54.
55. for node in graph:
56.     if node in forest:
57.         continue
58.     prim(node, graph, forest, forest_edges)
59.
60. for edge in forest_edges:
61.     print(f'{edge.first} - {edge.second}')