-
Notifications
You must be signed in to change notification settings - Fork 0
/
Prim.py
96 lines (78 loc) · 2.24 KB
/
Prim.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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
# Prim algorithm
# Find the minimum spanning tree of a connected undirect graph.
# Complexity O(n²)
import heapq
import operator
def prim(graph, start):
prev = {}
dist = {}
visited = set()
for node in graph:
prev[node] = -1
dist[node] = float('inf')
dist[start] = 0
remaining_nodes = [(dist[node], node) for node in graph]
while remaining_nodes:
heapq.heapify(remaining_nodes)
closest_node = heapq.heappop(remaining_nodes)[1]
visited.add(closest_node)
for node in get_neighbours(graph, closest_node):
if node not in visited:
if graph[closest_node][node] < dist[node]:
dist[node] = graph[closest_node][node]
prev[node] = closest_node
ind = list(map(operator.itemgetter(1), remaining_nodes)).index(node)
remaining_nodes[ind] = (dist[node], node)
return prev
def add_edge(graph, edge, weight):
edge = set(edge)
node1 = edge.pop()
if edge:
node2 = edge.pop()
else:
node2 = node1
graph[node1][node2] = weight
graph[node2][node1] = weight
def add_node(graph, node):
if node not in graph:
graph[node] = {}
def get_neighbours(graph, node):
nodes = []
if node in graph:
for neighbour in graph[node]:
if neighbour not in nodes:
nodes.append(neighbour)
return nodes
def edges_to_graph(edges):
graph = dict()
for edge in edges:
weight = edge.pop()
node1 = edge.pop()
if node1 not in graph:
add_node(graph, node1)
if edge:
node2 = edge.pop()
if node2 not in graph:
add_node(graph, node2)
else:
node2 = node1
graph[node2][node1] = weight
graph[node1][node2] = weight
return graph
def to_string(graph):
res = "Graph :\n"
for k in graph:
res += str(k) + " : " + str(graph[k]) + "\n"
return res
V = ['A', 'B', 'C', 'D', 'E', 'F']
E = [['A', 'B', 3],
['A', 'C', 5],
['B', 'C', 6],
['B', 'D', 7],
['C', 'D', 1],
['D', 'C', 2],
['E', 'F', 3],
['F', 'C', 5]]
H = edges_to_graph(E)
print(to_string(H))
print(prim(H,'A'))