CODEWITHSUNDEEP
pythonalgorithmminimum-spanning-treeadjacency-list-model
user1560400
user1560400

Reputation: 3

How to change this prim's algorithm function to take in and return an adjacency list?

I want to find the minimum spanning tree in a graph structured as an adjacency list. I was able to figure out the MST using Prim's algorithm but my current solution doesn't use an adjacency list.

from collections import defaultdict
from heapq import *

def prim( nodes, edges ):
   conn = defaultdict( list )
   for n1,n2,c in edges:
    conn[ n1 ].append( (c, n1, n2) )
    conn[ n2 ].append( (c, n2, n1) )

mst = []
used = set( nodes[ 0 ] )
usable_edges = conn[ nodes[0] ][:]
heapify( usable_edges )

while usable_edges:
    cost, n1, n2 = heappop( usable_edges )
    if n2 not in used:
        used.add( n2 )
        mst.append( ( n1, n2, cost ) )

        for e in conn[ n2 ]:
            if e[ 2 ] not in used:
                heappush( usable_edges, e )
 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 )

I need it to take and return something like this:

 {'A': [('B', 2)],
  'B': [('A', 2), ('C', 5)],
  'C': [('B', 5)]}

thank you!

Upvotes: 0

Views: 714

Answers (1)

kraskevich
kraskevich

Reputation: 18556

Once you have an mst as a list of edges, you just need to iterate over it to populate the dictionary.

That's exactly what you're doing in the very beginning of the prim function:

conn = defaultdict(list)
for n1, n2, c in edges:
    conn[n1].append((c, n1, n2))
    conn[n2].append((c, n2, n1))

You need to do exactly the same thing for the mst list instead of edges.

Upvotes: 0

Related Questions