-
Notifications
You must be signed in to change notification settings - Fork 0
/
graph.py
50 lines (38 loc) · 1.44 KB
/
graph.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
from pythonds.graphs import Graph, Vertex, PriorityQueue
import networkx as nx
import matplotlib.pyplot as plt
def minimum_spanning_tree(graph, initial_city):
start_city = graph.getVertex(initial_city)
heap = PriorityQueue()
aux_list = []
solution = []
visited_nodes = []
new_graph = Graph()
for conect in start_city.getConnections():
aux_list.append((start_city.getWeight(conect), (start_city.id, conect.id)))
visited_nodes.append(start_city.id)
heap.buildHeap(aux_list)
while not heap.isEmpty():
cur_pair = heap.delMin()
cur_node = graph.getVertex(cur_pair[1])
if cur_node.id not in visited_nodes:
solution.append(cur_pair)
visited_nodes.append(cur_node.id)
for conect in cur_node.getConnections():
if conect.id not in visited_nodes:
heap.add((cur_node.getWeight(conect), (cur_node.id, conect.id)))
else:
continue
print('Solution: ' + str(solution))
for city_1, city_2 in solution:
new_graph.addEdge(city_1, city_2, 0)
return new_graph
def show_graph(graph, graph_name):
visible_graph = nx.Graph()
for city in graph:
for conect in city.getConnections():
visible_graph.add_edge(city.id, conect.id)
visible_graph.add_node(city.id)
nx.draw_networkx(visible_graph, with_labels=True)
plt.savefig(graph_name)
plt.show()