Implementation of Prim's Algorithm to find minimum spanning tree

Question

What is the time complexity of the following code? I'm implementing Prim's Algorithm with adjacency matrix representation of the graph and priority queue.

from graph import adj_mtx, AP
import heapq as hq

lv, visited, h = float('inf'), {}, []  # O(n), lv stands for 'large_value' used to determine infinite distance, h is the heap 


def prims_mst(adj_matrix, src):
    hq.heappush(h, (0, (src, None)))  # O(logn)
    curr_dist = {item.value: lv if item.value != src else 0 for item in AP}  # curr_dist stands for current distances from the source

    while len(h) != 0:
        # print h[0]
        curr_nd = hq.heappop(h)[1][0]  # first element of the tuple is the value, second is the node  # O(1)
        visited[curr_nd] = True  # O(1)
        for nd, dst in enumerate(adj_matrix[src]):  # O(n) -> n is the number of nodes
            if nd not in visited and curr_dist[nd] > curr_dist[curr_nd] + adj_matrix[curr_nd][nd]:
                curr_dist[nd] = curr_dist[curr_nd] + adj_matrix[curr_nd][nd]
                hq.heappush(h, (curr_dist[nd], (nd, curr_nd)))  # O(logn)
        print h

Answer 1

1. Problem

The code in the post does not implement Prim's algorithm! Prim's algorithm constructs a minimum spanning tree. But the code in the post does not construct a tree.

(It could certainly be modified to construct the tree, but you have to actually do it!)

2. Analysis

The key thing to spot about the code in the post is that a node might get added to the queue more than once. This is becasue a node is added to the queue every time a new route is discovered to that node that is shorter than the previously known best route.

In particular, the \$n\$th node to be visited can be added to the queue up to \$n-1\$ times (because a new and better route might be discovered via any of the previous \$n-1\$ nodes to be visited). So the while loop is executed \$O(n^2)\$ times, and on each iteration of the while loop the inner for loop takes \$O(n)\$, giving a runtime of \$O(n^3)\$ overall.

But with an adjacency matrix it is possible to implement Prim's algorithm in \$O(n^2)\$ (see Wikipedia). So you have some work to do!

3. Other review points

1. There's no docstring. What does prims_mst do? What does it return? What arguments should I pass? I think that if you had stopped and written a docstring then you might have spotted the problem in section 1.

2. The variables visited and h are global. This means that you can only call prims_mst once — on the second call, the global variables have the wrong initial values. (Or else you have to remember to reset these variables, which is inconvenient and easy to forget.) It also means that the code can't be used from multiple threads.

3. Names like nd and lv and dst are hard to read. There is no danger of running out of vowels, so nothing will go wrong if you write node and distance instead.

4. Similarly, instead of importing a module and giving it a hard-to-read alias:

   import heapq as hq
   

   consider importing the names you need:

   from heapq import heappop, heappush
   

5. The variable lv is used only once, so you could just inline its value at the point of use.

6. Because the graph is given as an adjacency matrix (and because of the enumerate line), we know that the nodes can be labelled with consecutive integers. But the construction of curr_dist uses AP to get the nodes. It would be better to construct it in the way that you use it:

   # For each node, the shortest distance from src found so far.
   best_distance = [float('inf')] * len(adj_matrix)
   best_distance[src] = 0
   

   This avoids the dependency on AP.

7. The line:

   while len(h) != 0:
   

   can be simplified to:

   while h:
   

8. Instead of representing the set of visited nodes as a dictionary mapping nodes to True, use a set.

9. This line assigns a variable dst that is not used anywhere:

   for nd, dst in enumerate(adj_matrix[src]):
   

   An improvement would be to write:

   for node, distance in enumerate(adj_matrix[curr_nd]):
   

   and then you could use distance instead of adj_matrix[curr_nd][nd], saving a couple of indexing operations each time.

10. The heap entries are data structures of the form:

    (distance, (node, parent))
    

    The code for getting data out of these entries is hard to understand because the tuple lookup [1][0] doesn't have any distinctive meaning:

    curr_nd = hq.heappop(h)[1][0]
    

    It would be better to use a collections.namedtuple:

    HeapEntry = namedtuple('HeapEntry', 'distance node parent')
    

    and then you can write:

    current_node = heappop(h).node
    

    which is much easier to understand.