dsa.prim API documentation

# dsa.prim

Module to access functions for Prim's Algorithm.

View Source

     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

def prims_mst( graph, start: str, mst_graph=None) -> dsa.graph.AdjacencyListWeightedGraph: View Source

     6def prims_mst(graph, start: str, mst_graph=None) -> AdjacencyListWeightedGraph:
     7    """
     8    Returns an MST given a graph and starting vertex.
     9    (Future: return a Tree type instead of a Graph type)
    10
    11    Args:
    12        graph: The graph to search an MST from. (can be either an AdjacencyListWeightedGraph or AdjacencyMatrixWeightedGraph)
    13        start (string): The starting vertex label.
    14        mst_graph: An empty graph object to output the MST in to.
    15
    16    Returns:
    17        AdjacencyListWeightedGraph: the MST of the graph.
    18    """
    19    def add_adjacent(graph, pq: PriorityQueue, visited: set, node: str):
    20        """Add all adjacent vertices from the given node to the priority queue."""
    21        visited.add(node)
    22        for adjacent, weight in graph[node].items():
    23            if adjacent not in visited:
    24                pq.push(weight, (node, adjacent))  # Push edge with weight as priority
    25
    26    # todo: update this so that it will return the appropriate graph type
    27    if mst_graph is None:
    28        mst_graph = AdjacencyListWeightedGraph()
    29
    30    visited = set()
    31    pq = PriorityQueue()
    32    total_vertices = len(set(graph.vertices()))
    33
    34    add_adjacent(graph, pq, visited, start)
    35
    36    # While the priority queue is not empty and we haven't visited all vertices
    37    while not pq.is_empty() and len(visited) < total_vertices:
    38        weight, edge = pq.pop_pair()
    39        start, end = edge
    40        # If the end vertex has not been visited, add edge to the MST
    41        if end not in visited:
    42            mst_graph.add_edge(start, end, graph[start][end])
    43            # add adjacent vertices to the priority queue and mark the end vertex as visited
    44            add_adjacent(graph, pq, visited, end)
    45    return mst_graph

Returns an MST given a graph and starting vertex. (Future: return a Tree type instead of a Graph type)

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 mst_weight(graph) -> int: View Source

    47def mst_weight(graph) -> int:
    48    """
    49    Returns the total weight of a graph given a starting vertex
    50
    51    Args:
    52        graph: The graph to find the total edge weight of.
    53
    54    Returns:
    55        int: The total weight of the graph.
    56    """
    57    total_weight = 0
    58    visited = set()
    59    for start, end, weight in graph.edges():
    60        if (start, end) not in visited:
    61            total_weight += weight
    62            visited.add((start, end))
    63            visited.add((end, start))
    64    return total_weight

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.