A bit ago I covered Dijkstra’s Algorithm which I pointed out shares much in common with BFS because BFS is a special case of Dijkstra’s.

Relatedly, with just a very minor change to Dijkstra’s algorithm we can derive another useful algorithm to find the minimum spanning tree! This is called Prim’s Algorithm. I’m not going to go in depth on how Prim’s works1, but just how close it is to Dijkstra’s.

Both depend on the same Node class:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

class Node
  attr_reader :value
  attr_accessor :parent, :distance

  def initialize(value)
    @value = value
    @parent = nil
    @distance = Float::INFINITY
  end

  def to_s
    [value, distance]
  end
end

class WeightedNode < Node
  attr_accessor :edge_list
  def initialize(value)
    @edge_list = {}
    super
  end
end

Then we just need to set up our Nodes and put them into a graph. Here’s just a snippet of that plumbing (note we need to set the edge for both nodes):

1
2
3
4

one, two, three, four = Array.new(4) { |i| WeightedNode.new(i+1) }

one.edge_list[two] = 5
two.edge_list[one] = 5

Now we can play with graphs which are just Structs of Nodes and compare how different Prim’s and Dijkstra’s look.

As a refresher, here’s my version of Dijkstra’s:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

def dijkstra
  nodes.first.distance = 0
  pq = PQueue.new(nodes) { |a, b| a.distance < b.distance }

  until pq.empty?
    (current = pq.pop).edge_list.keys.each do |n|
      alt = current.distance + current.edge_list[n]
      if pq.include?(n) && alt < n.distance
        n.distance = alt
        n.parent = current
      end
    end
  end

  nodes.map(&:value).zip(nodes.map(&:distance))
end

And here’s my version of Prim’s:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

def prim
  nodes.first.distance = 0
  pq = PQueue.new(nodes) { |a, b| a.distance < b.distance }

  until pq.empty?
    (current = pq.pop).edge_list.keys.each do |n|
      alt = current.edge_list[n]
      if pq.include?(n) && alt < n.distance
        n.parent = current
        n.distance = alt
      end
    end
  end

  nodes.map(&:value).zip(nodes.map(&:distance))
end

Did you spot the difference?2

It’s a pretty subtle, and I hadn’t noticed the similarity of these two until I saw this SO answer. The “relaxing” function in Prim’s doesn’t accumulate. In a rough and informal sense, Dijkstra’s folds where Prim’s just needs to map.

So in Prim’s we have alt = current.edge_list[n] where in Dijkstra’s we have alt = current.distance + current.edge_list[n].

In concrete terms, that means we can rewrite this as one function, which operates on one of two procs:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

def dijkstra_proc
  proc do |c, n|
    c.distance + c.edge_list[n]
  end
end
def prim_proc
  proc do |c, n|
    c.edge_list[n]
  end
end

def both
  fail unless block_given?
  nodes.first.distance = 0
  pq = PQueue.new(nodes) { |a, b| a.distance < b.distance }

  until pq.empty?
    (current = pq.pop).edge_list.keys.each do |n|
      alt = yield(current, n)
      if pq.include?(n) && alt < n.distance
        n.distance = alt
        n.parent = current
      end
    end
  end

  nodes.map(&:value).zip(nodes.map(&:distance))
end

Pretty neat, I think! You can contrast this more “flexible” implementation to one designed specifically for Prim’s.

Doing Prim’s also alerted me to the downside of keeping distance on the Nodes as I’ve been doing. When I tried to run #dijkstra right after #prim I got the wrong answer, as both algorithms are dependent on the distances starting at Float::INFINITY. Side-effects! They’ll always bite you.

  1. There’s a lot of good explanations and demos, but I especially liked this one. ↩

  2. Check out line 7. ↩