Prim's Algorithm using heapq module in python

Question

Implementing Prim's with a \$O(m\log n)\$ run time using heaps, similar to Dijkstra heap implementation. Where \$n\$ is the number of nodes, and \$m\$ is the number of edges. Using the heapq library I tried to create a delete function for my heap ('unexplored') so that this heap only stores crossing edges at every iteration of outer While loop. That didn't work, I got amazing runtime but wrong result, so I just cut corners and used heapify and inf. But now the run time is around \$O(nm)\$, but I do get the correct answer now though.

graph is a dict, which contains adjacency lists of an undirected graph as:

{v1: [[v2, cost2], [v3, cost3],... ]... }

from math import inf
from heapq import heappush, heappop, _siftup, _siftdown, heapify
from time import time

def prims(graph):
    start = list(graph)[0]
    explored, treeCosts = set([start]), []

    unexplored = []
    for v in graph[start]:
        heappush(unexplored, [v[1], v[0]])

    while explored != set(graph):
        winner = heappop(unexplored)

        explored |= set([winner[1]])
        treeCosts.append(winner[0])

        #deleting edges which have head at 'winner'
        for v in unexplored:
            if v[1] == winner[1]:
                v[0] = inf
                '''i = unexplored.index(v)
                unexplored[i], unexplored[-1] = unexplored[-1], unexplored[i]
                unexplored.pop()
                if i < len(unexplored):
                    _siftup(unexplored, i)
                    _siftdown(unexplored, 0, i)'''
        heapify(unexplored)

        #adding the new crossing edges to heap
        for v in graph[winner[1]]:
            if v[0] not in explored:
                heappush(unexplored, [v[1], v[0]])

    return sum(treeCosts)

if __name__ == '__main__':
    draft, graph = open('undirectedGraph (weighted, 500).txt').read().splitlines(), {}

    for line in draft[1:]:
        edge = list(map(int, line.split()))
        if edge[0] not in graph:
            graph[edge[0]] = []
        if edge[1] not in graph:
            graph[edge[1]] = []

        graph[edge[0]].append(edge[1:])
        graph[edge[1]].append([edge[0], edge[2]])

    startTime = time()
    print('Good prims: ' + str(prims(graph)) + ', ' + 'Time: ' + str(time() - startTime))

Answer 1

1. Review

1. There is no docstring. What does prims do? What argument does it take? What does it return? The text in the post would be a good starting point for a docstring.

2. The code has a lot of sequence lookups [1] and [0]. This makes it hard for the reader to follow (what does [1] mean?). It's clearer to use Python's sequence unpacking to provide names for these elements. For example, instead of:

   for v in graph[start]:
       heappush(unexplored, [v[1], v[0]])
   

   write something like:

   for neighbour, cost in graph[start]:
       heappush(unexplored, (cost, neighbour))
   

   This makes the data structures easier for the reader to follow.

3. To pick an arbitrary vertex from graph, the code makes a list of all the vertices, which is then thrown away because only the first vertex in the list is needed:

   start = list(graph)[0]
   

   It would be better not to make the list at all:

   start = next(iter(graph))   # Arbitrary starting vertex
   

4. Instead of keeping a list of costs and summing them at the end:

   treeCosts = []
   # ...
   treeCosts.append(winner[0])
   # ...
   return sum(treeCosts)
   

   it would be simpler to maintain a running total:

   total = 0
   # ...
   total += winner[0]
   # ...
   return total
   

5. Every iteration, the code tests that all vertices have been explored:

   while explored != set(graph):
   

   This takes time proportional to the number of vertices in the graph. It would be better to test whether there are any unexplored edges remaining:

   while unexplored:
   

   which can be tested in constant time.

6. Once a winning vertex has been found and added to the tree, all other edges in the heap that end at that vertex are now invalid. The code tries to delete them from the heap by setting their cost to infinity, but this violates the heap invariant. I can see that you attempted to write some code for trying to restore the heap invariant, but this didn't work and you ended up calling heapify instead, but this takes linear time and so leads to unacceptable (quadratic) overall runtime complexity.

   This is a problem that often arises when using a heap: how do you efficiently delete an item? Well, there is a way to do it, but it is somewhat complicated (see this answer on Stack Overflow for a description of the method).

   However, there is a much simpler solution to this problem: leave the bad item in the heap and postpone dealing with it until it is popped. (See the fourth bullet under "Priority Queue Implementation Notes" in the heapq documentation.) In the minimum spanning tree algorithm this is easy because you can just check against the explored set, like this:

   while unexplored:
       cost, winner = heappop(unexplored)
       if winner in explored:
           continue # Already found via a shorter path
   

   You should check that leaving these items in the heap does not affect the runtime complexity (hint: \$O(\log {n^2}) = O(\log n)\$).

7. The heap is initialized with all the neighbours of the starting vertex:

   unexplored = []
   for v in graph[start]:
       heappush(unexplored, [v[1], v[0]])
   

   But this duplicates the code for adding the neighbours of the winning vertex:

   #adding the new crossing edges to heap
   for v in graph[winner[1]]:
       if v[0] not in explored:
           heappush(unexplored, [v[1], v[0]])
   

   It would be simpler to start with just the starting vertex in the heap (at zero cost), and the explored set empty, like this:

   explored = set()
   unexplored = [(0, start)]
   

2. Revised code

from heapq import heappush, heappop

def minimum_spanning_tree_cost(graph):
    """Return the sum of the costs of the edges in the minimum spanning
    tree for the given graph, which must be a mapping from nodes to an
    iterable of (neighbour, edge-cost) pairs.

    """
    total = 0                   # Total cost of edges in tree
    explored = set()            # Set of vertices in tree
    start = next(iter(graph))   # Arbitrary starting vertex
    unexplored = [(0, start)]   # Unexplored edges ordered by cost
    while unexplored:
        cost, winner = heappop(unexplored)
        if winner not in explored:
            explored.add(winner)
            total += cost
            for neighbour, cost in graph[winner]:
                if neighbour not in explored:
                    heappush(unexplored, (cost, neighbour))
    return total

3. What if vertices can't be compared?

The heap contains pairs (cost, vertex) which get compared by the heap algorithm to see which is smallest. If two cost values are the same then the heap algorithm will have to compare the vertices to break the tie. This is fine if vertices can be compared (for example if vertices are represented by numbers or strings) but if they cannot be compared then the code will fail with:

TypeError: unorderable types

This is a common problem: see the second bullet under "Priority Queue Implementation Notes" in the heapq documentation. Again, there is a general solution: add a tie-breaking value between the cost and the vertex. itertools.count is a good source of tie-breaking values.

from heapq import heappush, heappop
from itertools import count

def minimum_spanning_tree_cost(graph):
    """Return the sum of the costs of the edges in the minimum spanning
    tree for the given graph, which must be a mapping from nodes to an
    iterable of (neighbour, edge-cost) pairs.

    """
    tiebreak = count().__next__ # Factory for tie-breaking values
    total = 0                   # Total cost of edges in tree
    explored = set()            # Set of vertices in tree
    start = next(iter(graph))   # Arbitrary starting vertex
    unexplored = [(0, tiebreak(), start)] # Unexplored edges ordered by cost
    while unexplored:
        cost, _, winner = heappop(unexplored)
        if winner not in explored:
            explored.add(winner)
            total += cost
            for neighbour, cost in graph[winner]:
                if neighbour not in explored:
                    heappush(unexplored, (cost, tiebreak(), neighbour))
    return total