A class representing an undirected weighted graph.
| 37 | |
| 38 | |
| 39 | class Graph: |
| 40 | """ |
| 41 | A class representing an undirected weighted graph. |
| 42 | """ |
| 43 | |
| 44 | def __init__(self, vertices: set[int], edges: Mapping[EdgeT, int]) -> None: |
| 45 | self.vertices: set[int] = vertices |
| 46 | self.edges: dict[EdgeT, int] = { |
| 47 | (min(edge), max(edge)): weight for edge, weight in edges.items() |
| 48 | } |
| 49 | |
| 50 | def add_edge(self, edge: EdgeT, weight: int) -> None: |
| 51 | """ |
| 52 | Add a new edge to the graph. |
| 53 | >>> graph = Graph({1, 2}, {(2, 1): 4}) |
| 54 | >>> graph.add_edge((3, 1), 5) |
| 55 | >>> sorted(graph.vertices) |
| 56 | [1, 2, 3] |
| 57 | >>> sorted([(v,k) for k,v in graph.edges.items()]) |
| 58 | [(4, (1, 2)), (5, (1, 3))] |
| 59 | """ |
| 60 | self.vertices.add(edge[0]) |
| 61 | self.vertices.add(edge[1]) |
| 62 | self.edges[(min(edge), max(edge))] = weight |
| 63 | |
| 64 | def prims_algorithm(self) -> Graph: |
| 65 | """ |
| 66 | Run Prim's algorithm to find the minimum spanning tree. |
| 67 | Reference: https://en.wikipedia.org/wiki/Prim%27s_algorithm |
| 68 | >>> graph = Graph({1,2,3,4},{(1,2):5, (1,3):10, (1,4):20, (2,4):30, (3,4):1}) |
| 69 | >>> mst = graph.prims_algorithm() |
| 70 | >>> sorted(mst.vertices) |
| 71 | [1, 2, 3, 4] |
| 72 | >>> sorted(mst.edges) |
| 73 | [(1, 2), (1, 3), (3, 4)] |
| 74 | """ |
| 75 | subgraph: Graph = Graph({min(self.vertices)}, {}) |
| 76 | min_edge: EdgeT |
| 77 | min_weight: int |
| 78 | edge: EdgeT |
| 79 | weight: int |
| 80 | |
| 81 | while len(subgraph.vertices) < len(self.vertices): |
| 82 | min_weight = max(self.edges.values()) + 1 |
| 83 | for edge, weight in self.edges.items(): |
| 84 | if (edge[0] in subgraph.vertices) ^ ( |
| 85 | edge[1] in subgraph.vertices |
| 86 | ) and weight < min_weight: |
| 87 | min_edge = edge |
| 88 | min_weight = weight |
| 89 | |
| 90 | subgraph.add_edge(min_edge, min_weight) |
| 91 | |
| 92 | return subgraph |
| 93 | |
| 94 | |
| 95 | def solution(filename: str = "p107_network.txt") -> int: |
no outgoing calls
no test coverage detected