diff --git a/greedy_methods/Prim's_algorithm.py b/greedy_methods/Prim's_algorithm.py new file mode 100644 index 000000000000..eaae76f6611f --- /dev/null +++ b/greedy_methods/Prim's_algorithm.py @@ -0,0 +1,49 @@ +from typing import Dict, Any +import heapq + + +def prim_mst(graph: Dict[Any, Dict[Any, int]]) -> Dict[Any, Any]: + """ + Generate the Minimum Spanning Tree (MST) using Prim's algorithm. + + Args: + graph (Dict[Any, Dict[Any, int]]): Adjacency dictionary of the graph where keys are nodes and + values are dictionaries of neighbor nodes with edge weights. + + Returns: + Dict[Any, Any]: A dictionary representing the MST with each node pointing to its parent in the MST. + + Example: + >>> graph = { + ... 'A': {'B': 2, 'C': 3}, + ... 'B': {'A': 2, 'C': 1}, + ... 'C': {'A': 3, 'B': 1} + ... } + >>> prim_mst(graph) + {'A': 'B', 'B': 'C', 'C': 'B'} + """ + if not graph: + return {} + + start_node = next(iter(graph)) + visited = set([start_node]) + edges = [(weight, start_node, to) for to, weight in graph[start_node].items()] + heapq.heapify(edges) + mst = {} + + while edges: + weight, frm, to = heapq.heappop(edges) + if to not in visited: + visited.add(to) + mst[to] = frm + for next_to, next_weight in graph[to].items(): + if next_to not in visited: + heapq.heappush(edges, (next_weight, to, next_to)) + + return mst + + +if __name__ == "__main__": + import doctest + + doctest.testmod()