perf(prim): index the edges by endpoint instead of rescanning them

src/undirected/prim.rs

@@ -1,12 +1,16 @@
 //! Find minimum-spanning-tree in an undirected graph using [Prim's
 //! algorithm](https://en.wikipedia.org/wiki/Prim%27s_algorithm).
 
+use rustc_hash::{FxHashMap, FxHashSet};
 use std::cmp::Reverse;
-use std::collections::{BinaryHeap, HashSet};
+use std::collections::BinaryHeap;
 use std::hash::Hash;
 
 /// Find a minimum-spanning-tree. From a collection of weighted edges,
 /// return a vector of edges forming a minimum-spanning-tree.
+///
+/// Edges are undirected: `(a, b, c)` and `(b, a, c)` describe the same edge, and either form
+/// may be used. The tree is grown from the first endpoint of the first edge.
 pub fn prim<N, C>(edges: &[(N, N, C)]) -> Vec<(&N, &N, C)>
 where
     N: Hash + Eq + Ord,
@@ -16,39 +20,41 @@ where
         return vec![];
     };
 
-    // Edges are undirected, so an edge touching the starting node may name it either first
-    // or second. Both forms have to be offered here: once the loop below has marked the
-    // starting node as visited, neither of its checks can reach back to it.
-    let mut priority_queue = edges
-        .iter()
-        .filter_map(|(n, n1, c)| {
-            if n == start {
-                Some(Reverse((c, n, n1)))
-            } else if n1 == start {
-                Some(Reverse((c, n1, n)))
-            } else {
-                None
-            }
-        })
-        .collect::<BinaryHeap<_>>();
+    // Index every edge under both of its endpoints, once. Growing the tree then only looks at
+    // the edges leaving the node just added, instead of rescanning the whole edge list for
+    // each of them.
+    let mut incident: FxHashMap<&N, Vec<(&C, &N)>> = FxHashMap::default();
+    for (a, b, cost) in edges {
+        incident.entry(a).or_default().push((cost, b));
+        if a != b {
+            incident.entry(b).or_default().push((cost, a));
+        }
+    }
 
-    let (mut mst, mut visited) = (Vec::new(), HashSet::new());
+    let mut mst = Vec::new();
+    let mut visited: FxHashSet<&N> = FxHashSet::default();
     visited.insert(start);
-    while let Some(Reverse((c, n, n1))) = priority_queue.pop() {
-        if visited.contains(n1) {
-            continue;
+    let mut priority_queue = BinaryHeap::new();
+    let mut grown = Some(start);
+    while let Some(node) = grown.take() {
+        // Offer every edge that leaves the tree through the node just added...
+        if let Some(candidates) = incident.get(node) {
+            for &(cost, other) in candidates {
+                if !visited.contains(other) {
+                    priority_queue.push(Reverse((cost, node, other)));
+                }
+            }
         }
-
-        mst.push((n, n1, c.clone()));
-
-        for (n2, n3, c) in edges {
-            if n1 == n2 && !visited.contains(n3) {
-                priority_queue.push(Reverse((c, n1, n3)));
-            } else if n1 == n3 && !visited.contains(n2) {
-                priority_queue.push(Reverse((c, n1, n2)));
+        // ... then take the cheapest edge that still reaches a new node.
+        while let Some(Reverse((cost, from, to))) = priority_queue.pop() {
+            if visited.contains(to) {
+                continue;
             }
+            mst.push((from, to, cost.clone()));
+            visited.insert(to);
+            grown = Some(to);
+            break;
         }
-        visited.insert(n1);
     }
     mst
 }