dsa.prim

 1""" Module to access functions for Prim's Algorithm. """
 2from dsa.graph import AdjacencyListWeightedGraph
 3from dsa.heap import PriorityQueue
 4    
 5def prims_mst(graph, start: str, mst_graph=None) -> AdjacencyListWeightedGraph:
 6    """
 7    Returns an MST given a graph and starting vertex.
 8    (Future: return a Tree type instead of a Graph type)
 9
10    Args:
11        graph: The graph to search an MST from. (can be either an AdjacencyListWeightedGraph or AdjacencyMatrixWeightedGraph)
12        start (string): The starting vertex label.
13        mst_graph: An empty graph object to output the MST in to.
14
15    Returns:
16        AdjacencyListWeightedGraph: the MST of the graph.
17    """
18    def add_adjacent(graph, pq: PriorityQueue, visited: set, node: str):
19        """Add all adjacent vertices from the given node to the priority queue."""
20        visited.add(node)
21        for adjacent, weight in graph[node].items():
22            if adjacent not in visited:
23                pq.push(weight, (node, adjacent))  # Push edge with weight as priority
24
25    # todo: update this so that it will return the appropriate graph type
26    if mst_graph is None:
27        mst_graph = AdjacencyListWeightedGraph()
28
29    visited = set()
30    pq = PriorityQueue()
31    total_vertices = len(set(graph.vertices()))
32
33    add_adjacent(graph, pq, visited, start)
34
35    # While the priority queue is not empty and we haven't visited all vertices
36    while not pq.is_empty() and len(visited) < total_vertices:
37        weight, edge = pq.pop_pair()
38        start, end = edge
39        # If the end vertex has not been visited, add edge to the MST
40        if end not in visited:
41            mst_graph.add_edge(start, end, graph[start][end])
42            # add adjacent vertices to the priority queue and mark the end vertex as visited
43            add_adjacent(graph, pq, visited, end)
44    return mst_graph
45
46def mst_weight(graph) -> int:
47    """
48    Returns the total weight of a graph given a starting vertex
49    
50    Args:
51        graph: The graph to find the total edge weight of.
52
53    Returns:
54        int: The total weight of the graph.
55    """
56    total_weight = 0
57    visited = set()
58    for start, end, weight in graph.edges():
59        if (start, end) not in visited:
60            total_weight += weight
61            visited.add((start, end))
62            visited.add((end, start))
63    return total_weight