Solve problem 107 with Prim's algorithm in Python

python/e107.py

@@ -0,0 +1,89 @@
+from dataclasses import dataclass
+from typing import List
+
+
+@dataclass
+class Edge:
+    weight: int
+    source_id: int
+    target_id: int
+
+
+@dataclass
+class Vertex:
+    id: int
+    edges: List[Edge]
+
+
+def load_matrix(filename):
+    def line_to_row(line):
+        r = []
+        for field in line.strip().split(","):
+            if field == "-":
+                r.append(0)
+            else:
+                r.append(int(field))
+        return r
+
+    with open(filename, "r") as f:
+        m = [line_to_row(l) for l in f]
+    return m
+
+
+def matrix_to_graph(matrix):
+    n_nodes = len(matrix)
+    graph = []
+    for row_index in range(n_nodes):
+
+        edges = []
+        for col_index in range(n_nodes):
+            weight = matrix[row_index][col_index]
+            if weight != 0:
+                edges.append(Edge(weight, row_index, col_index))
+        graph.append(Vertex(row_index, edges))
+    return graph
+
+
+def total_weight(graph):
+    total_weight = 0
+    for vertex in graph:
+        for edge in vertex.edges:
+            total_weight += edge.weight
+    total_weight //= 2
+    return total_weight
+
+
+def prims_algo(graph):
+    """ Takes a graph and returns the edges that for a minimum spanning tree
+    for that graph. """
+    visited_id = set([0])
+    min_edges = []
+    while len(visited_id) != len(graph):
+        min_edge = None
+        for vertex_id in visited_id:
+            vertex = graph[vertex_id]
+            for edge in vertex.edges:
+                if edge.target_id in visited_id:
+                    pass
+                elif min_edge is None:
+                    min_edge = edge
+                elif edge.weight < min_edge.weight:
+                    min_edge = edge
+        visited_id.add(min_edge.target_id)
+        min_edges.append(min_edge)
+    return min_edges
+
+
+def euler_107():
+    m = load_matrix("../txt/e107.txt")
+    g = matrix_to_graph(m)
+    weight = total_weight(g)
+    new_weight = sum([e.weight for e in prims_algo(g)])
+    return weight - new_weight
+
+
+if __name__ == "__main__":
+    solution = euler_107()
+    print("e107.py: " + str(solution))
+    assert(solution == 259679)
+