66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
import math
|
|
import re
|
|
from graph import Graph, AdjacencyListGraph, AdjacencyMatrixGraph, NodeColor, Vertex
|
|
import heapq
|
|
|
|
|
|
def mst_prim(self, start_name: str = None):
|
|
|
|
distance_map = {}
|
|
parent_map = {}
|
|
|
|
Vertex.__lt__ = lambda self, other: distance_map[self] < distance_map[other]
|
|
|
|
queue = []
|
|
|
|
if start_name is None:
|
|
start_name = self.all_vertices()[0].value;
|
|
|
|
# Initialize the maps
|
|
for vertex in self.all_vertices():
|
|
distance_map[vertex] = 0 if vertex.value == start_name else math.inf
|
|
parent_map[vertex] = None
|
|
heapq.heappush(queue, vertex)
|
|
|
|
# Process the queue
|
|
cost = 0
|
|
while len(queue) > 0:
|
|
vertex = heapq.heappop(queue)
|
|
cost += distance_map[vertex]
|
|
for (dest, w) in self.get_adjacent_vertices_with_weight(vertex.value):
|
|
if dest in queue and distance_map[dest] > w:
|
|
queue.remove(dest)
|
|
distance_map[dest] = w
|
|
parent_map[dest] = vertex
|
|
heapq.heappush(queue, dest)
|
|
|
|
# Return the distance and predecessor maps
|
|
return parent_map, cost
|
|
|
|
AdjacencyListGraph.mst_prim = mst_prim
|
|
AdjacencyMatrixGraph.mst_prim = mst_prim
|
|
|
|
|
|
def read_elektro_into_graph(graph: Graph, filename: str):
|
|
pattern = re.compile(r'"([^"]+)";"([^"]+)";(\d+)')
|
|
with (open(filename, "r") as file):
|
|
for line in file:
|
|
m = pattern.match(line)
|
|
if m:
|
|
start_name = m.group(1)
|
|
end_name = m.group(2)
|
|
cost = int(m.group(3))
|
|
graph.insert_vertex(start_name)
|
|
graph.insert_vertex(end_name)
|
|
graph.connect(start_name, end_name, cost)
|
|
graph.connect(end_name, start_name, cost)
|
|
|
|
if __name__ == "__main__":
|
|
|
|
graph = AdjacencyMatrixGraph()
|
|
read_elektro_into_graph(graph, "../../elektro.txt")
|
|
parents, cost = graph.mst_prim()
|
|
print(f"Kosten {cost}")
|
|
for node, parent in parents.items():
|
|
if parent is not None:
|
|
print(f"{node} - {parent}") |