graph_algorithms.py 文件源码

def minimum_spanning_tree(gr):
    """ Uses prim's algorithm to return the minimum 
    cost spanning tree in a undirected connected graph.
    Works only with undirected and connected graphs """
    s = gr.nodes()[0] 
    nodes_explored = set([s])
    nodes_unexplored = gr.nodes()
    nodes_unexplored.remove(s)
    min_cost, node_heap = 0, []

    #computes the key for each vertex in unexplored
    for n in nodes_unexplored:
        min = compute_key(gr, n, nodes_explored)
        heapq.heappush(node_heap, (min, n))

    while len(nodes_unexplored) > 0:
        # adds the cheapest to "explored"
        node_cost, min_node = heapq.heappop(node_heap)
        min_cost += node_cost
        nodes_explored.add(min_node)
        nodes_unexplored.remove(min_node)

        # recompute keys for neighbors of deleted node
        for v in gr.neighbors(min_node):
            if v in nodes_unexplored:
                for i in range(len(node_heap)):
                    if node_heap[i][1] == v:
                        node_heap[i] = (compute_key(gr, v, nodes_explored), v)
                        heapq.heapify(node_heap)
    return min_cost