prim for MST

mst_prim.py

def prim(nodes, edges): 
    conn = {}
    for u, v, w in edges:
        if u not in conn:
            conn[u] = [(w,u,v)]            
        else:
            conn[u].append((w, u, v))
        if v not in conn:
            conn[v] = [(w,v,u)]
        else:
            conn[v].append((w, v, u))            
    node = nodes[0]    
    usable_edges = conn[node]    
    heapq.heapify(usable_edges)
    Q = set(node)
    mst = []
    while len(usable_edges) > 0:        
        wt, from_node, to_node = heapq.heappop(usable_edges)        
        if to_node not in Q:
            Q.add(to_node)
            mst.append((from_node, to_node, wt))            
            for edge in conn[to_node]:                
                ww, uu, vv = edge
                if vv not in Q:
                    heapq.heappush(usable_edges, (ww, uu, vv))        
    return mst


#test
nodes = list("ABCDEFG")
edges = [ ("A", "B", 7), ("A", "D", 5),
          ("B", "C", 8), ("B", "D", 9), ("B", "E", 7),
      ("C", "E", 5),
      ("D", "E", 15), ("D", "F", 6),
      ("E", "F", 8), ("E", "G", 9),
      ("F", "G", 11)]
 
print("prim:", prim( nodes, edges ))