-
Notifications
You must be signed in to change notification settings - Fork 0
/
Dijkstra.py
92 lines (72 loc) · 3.02 KB
/
Dijkstra.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
import heapq
class Graph:
def __init__(self):
self.adjacency_list = {}
def add_edge(self, u, v, weight):
if u not in self.adjacency_list:
self.adjacency_list[u] = []
if v not in self.adjacency_list:
self.adjacency_list[v] = []
self.adjacency_list[u].append((v, weight))
self.adjacency_list[v].append((u, weight))
def dijkstra_mst(self, source):
min_heap = [] # Priority queue to store vertices based on their minimum distance
heapq.heappush(min_heap, (0, source)) # Push the source vertex with distance 0
mst_cost = 0 # Total cost of the MST
mst_edges = [] # List to store MST edges
# Dictionary to keep track of minimum distance to each vertex
min_distance = {vertex: float('inf') for vertex in self.adjacency_list}
min_distance[source] = 0
# Set to track visited vertices
visited = set()
while min_heap:
current_distance, current_vertex = heapq.heappop(min_heap)
if current_vertex in visited:
continue
# Mark the current vertex as visited
visited.add(current_vertex)
mst_cost += current_distance
if current_vertex != source:
mst_edges.append((min_distance[current_vertex], current_vertex))
# Explore neighbors of the current vertex
for neighbor, weight in self.adjacency_list[current_vertex]:
if neighbor not in visited and weight < min_distance[neighbor]:
min_distance[neighbor] = weight
heapq.heappush(min_heap, (weight, neighbor))
return mst_cost, mst_edges
def get_input_edges():
edges = []
while True:
print("Enter an edge (format: u v weight) or enter 'done' to finish:")
user_input = input().strip()
if user_input.lower() == 'done':
break
try:
u, v, weight = user_input.split()
weight = int(weight)
edges.append((u, v, weight))
except ValueError:
print("Invalid input format. Please try again.")
return edges
def main():
# Create a new graph instance
graph = Graph()
# Prompt the user to enter edges and weights
edges = get_input_edges()
# Add edges to the graph
for u, v, weight in edges:
graph.add_edge(u, v, weight)
# Specify the source vertex for MST
source_vertex = input("Enter the source vertex for MST: ").strip()
# Find the Minimal Spanning Tree (MST) using Dijkstra's algorithm
if source_vertex in graph.adjacency_list:
mst_cost, mst_edges = graph.dijkstra_mst(source_vertex)
# Display the Minimal Spanning Tree (MST) cost and edges
print("Minimal Spanning Tree (MST) Cost:", mst_cost)
print("Minimal Spanning Tree (MST) Edges:")
for cost, edge in mst_edges:
print(edge)
else:
print("Invalid source vertex. Please enter a valid vertex.")
if __name__ == "__main__":
main()