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)