Module dsa.prim

Module to access functions for Prim's Algorithm.

Expand source code
""" Module to access functions for Prim's Algorithm. """
from dsa.graph import AdjacencyListWeightedGraph
from dsa.heap import PriorityQueue
    
def prims_mst(graph, start: str, mst_graph=None) -> AdjacencyListWeightedGraph:
    """
    Returns an MST given a graph and starting vertex.

    Args:
        graph: The graph to search an MST from. (can be either an AdjacencyListWeightedGraph or AdjacencyMatrixWeightedGraph)
        start (string): The starting vertex label.
        mst_graph: an empty graph object to output the MST in to 

    Returns:
        AdjacencyListWeightedGraph: the MST of the graph.
    """
    def add_adjacent(graph, pq: PriorityQueue, visited: set, node: str):
        """Add all outgoing edges from the given node to the priority queue."""
        visited.add(node)
        for neighbor, weight in graph[node].items():
            if neighbor not in visited:
                pq.push(weight, (node, neighbor))  # Push edge with weight as priority

    # update this so that it will return the appropriate graph type
    if mst_graph is None:
        mst_graph = AdjacencyListWeightedGraph()

    visited = set()
    pq = PriorityQueue()

    add_adjacent(graph, pq, visited, start)
    vertices_set = set(graph.vertices())

    while not pq.is_empty() and visited != vertices_set:
        weight, edge = pq.pop_pair()
        u, v = edge
        if v not in visited:
            mst_graph.add_edge(u, v, graph[u][v])
            add_adjacent(graph, pq, visited, v)
    return mst_graph

def mst_weight(graph) -> int:
    """
    Returns the total weight of a graph given a starting vertex
    
    Args:
        graph: The graph to find the total edge weight of.

    Returns:
        int: The total weight of the graph.
    """
    total_weight = 0
    visited = set()
    for start, end, weight in graph.edges():
        if (start, end) not in visited:
            total_weight += weight
            visited.add((start, end))
            visited.add((end, start))
    return total_weight

Functions

def mst_weight(graph) ‑> int

Returns the total weight of a graph given a starting vertex

Args

graph
The graph to find the total edge weight of.

Returns

int
The total weight of the graph.
Expand source code
def mst_weight(graph) -> int:
    """
    Returns the total weight of a graph given a starting vertex
    
    Args:
        graph: The graph to find the total edge weight of.

    Returns:
        int: The total weight of the graph.
    """
    total_weight = 0
    visited = set()
    for start, end, weight in graph.edges():
        if (start, end) not in visited:
            total_weight += weight
            visited.add((start, end))
            visited.add((end, start))
    return total_weight
def prims_mst(graph, start: str, mst_graph=None) ‑> AdjacencyListWeightedGraph

Returns an MST given a graph and starting vertex.

Args

graph
The graph to search an MST from. (can be either an AdjacencyListWeightedGraph or AdjacencyMatrixWeightedGraph)
start : string
The starting vertex label.
mst_graph
an empty graph object to output the MST in to

Returns

AdjacencyListWeightedGraph
the MST of the graph.
Expand source code
def prims_mst(graph, start: str, mst_graph=None) -> AdjacencyListWeightedGraph:
    """
    Returns an MST given a graph and starting vertex.

    Args:
        graph: The graph to search an MST from. (can be either an AdjacencyListWeightedGraph or AdjacencyMatrixWeightedGraph)
        start (string): The starting vertex label.
        mst_graph: an empty graph object to output the MST in to 

    Returns:
        AdjacencyListWeightedGraph: the MST of the graph.
    """
    def add_adjacent(graph, pq: PriorityQueue, visited: set, node: str):
        """Add all outgoing edges from the given node to the priority queue."""
        visited.add(node)
        for neighbor, weight in graph[node].items():
            if neighbor not in visited:
                pq.push(weight, (node, neighbor))  # Push edge with weight as priority

    # update this so that it will return the appropriate graph type
    if mst_graph is None:
        mst_graph = AdjacencyListWeightedGraph()

    visited = set()
    pq = PriorityQueue()

    add_adjacent(graph, pq, visited, start)
    vertices_set = set(graph.vertices())

    while not pq.is_empty() and visited != vertices_set:
        weight, edge = pq.pop_pair()
        u, v = edge
        if v not in visited:
            mst_graph.add_edge(u, v, graph[u][v])
            add_adjacent(graph, pq, visited, v)
    return mst_graph