def mst_kruskal(graph):
"""
Returns a minimum spanning tree of the input graph using Kruskal's algorithm
"""
edge_set = graph.edge_set
sortedEdges = sorted(edge_set, key = operator.attrgetter('weight'))
component_membership = {vertex:UpTree() for vertex in set(itertools.chain.from_iterable(edge_set))}
mst_edges = set()
for edge in sortedEdges:
ep_1, ep_2 = edge
if component_membership[ep_1].find() != component_membership[ep_2].find():
component_membership[ep_1].union(component_membership[ep_2])
mst_edges.add(edge)
if(len(mst_edges) == graph.size - 1):
break
#edges will now be referencing vertices with different object ids than those vertices in the vertexSet
mst_edges = copy.deepcopy(mst_edges)
mst_vertices = set(itertools.chain.from_iterable(mst_edges))
for vertex in mst_vertices:
vertex._neighbors.clear()
return Graph._from_vertex_and_edge_sets(vertex_set = mst_vertices, edge_set = mst_edges)
def mst_prim(graph):
"""
Returns a minimum spanning tree (MST) of the input graph using Prim's algorithm
"""
mst_edges = set()
unincluded_vertices = graph.vertex_set
optimal_connecting_edge = {vertex: None for vertex in unincluded_vertices}
optimal_connecting_distance = {
vertex: sys.maxsize
for vertex in unincluded_vertices
}
initial_vertex = unincluded_vertices.pop()
#assumes no parallel edges
for edge in graph._outgoing_edges(initial_vertex):
adjacent_vertex = edge.get_adjacent_vertex(initial_vertex)
if edge.weight < optimal_connecting_distance[adjacent_vertex]:
optimal_connecting_edge[adjacent_vertex] = edge
optimal_connecting_distance[adjacent_vertex] = edge.weight
while unincluded_vertices:
new_vertex = min(unincluded_vertices,
key=optimal_connecting_distance.get)
new_edge = optimal_connecting_edge.get(new_vertex)
mst_edges.add(new_edge)
for edge in graph._outgoing_edges(new_vertex):
adjacent_vertex = edge.get_adjacent_vertex(new_vertex)
if edge.weight < optimal_connecting_distance[adjacent_vertex]:
optimal_connecting_edge[adjacent_vertex] = edge
optimal_connecting_distance[adjacent_vertex] = edge.weight
unincluded_vertices.remove(new_vertex)
mst_edges = copy.deepcopy(mst_edges)
mst_vertices = set(itertools.chain.from_iterable(mst_edges))
#still need to clear old neighbors from vertices
for vertex in mst_vertices:
vertex._neighbors.clear()
return Graph._from_vertex_and_edge_sets(vertex_set=mst_vertices,
edge_set=mst_edges)