def prim(graph: UndirectedGraph) -> UndirectedGraph:
"""
Prim's algorithm is a greedy algorithm that finds a minimum spanning tree for a weighted undirected graph.
This means it finds a subset of the edges that forms a tree that includes every vertex
O(E + V log V)
"""
current_vertex = graph.get_vertices()[0]
new_graph = UndirectedGraph()
heap = FibonacciHeap(map(lambda x: (x, current_vertex), graph.get_outgoing_edges(current_vertex)))
vertices = {current_vertex}
while not heap.is_empty():
edge, current_vertex = heap.pop()
next_vertex = edge.a if edge.a != current_vertex else edge.b
if next_vertex in vertices:
continue
# Add edge to the minimum spanning tree
vertices.add(next_vertex)
new_graph.add_edge(current_vertex, next_vertex, edge.weight)
# Add all edges from this new point to the heap
[heap.insert((edge, next_vertex)) for edge in graph.get_outgoing_edges(next_vertex)]
return new_graph
