Add power connection and steiner tree

factorio-graph/src/wheighted_graph/steiner_tree.rs

@@ -2,16 +2,72 @@ use crate::priority_queue::PriorityQueue;
 use std::hash::Hash;
 use std::{collections::HashSet, fmt::Debug};
 
+use super::shortest_path::dijkstra;
 use super::{WheightedGraph, shortest_path::QueueObject};
 
-pub fn takaheshi_matsuyama<G, P>(graph: G, nodes: &[G::Node]) -> Option<Vec<Vec<G::Node>>>
+struct MultistartWrapper<'a, 'b, G>
 where
-    P: PriorityQueue<QueueObject<G::Node>> + Debug,
-    P::Handle: Debug,
-    G::Node: Eq + Hash + Clone,
     G: WheightedGraph,
 {
-    let end_nodes: HashSet<G::Node> = HashSet::from_iter(nodes.iter().cloned());
-
-    todo!()
+    graph: &'a G,
+    start_nodes: &'b [G::Node],
+}
+
+impl<G> WheightedGraph for MultistartWrapper<'_, '_, G>
+where
+    G: WheightedGraph,
+    G::Node: Clone,
+{
+    type Node = Option<G::Node>;
+
+    fn edge(&self, node: &Self::Node, num: usize) -> Option<(Self::Node, f64)> {
+        match node {
+            Some(n) => self.graph.edge(n, num).map(|(n, v)| (Some(n), v)),
+            None => self.start_nodes.get(num).cloned().map(|n| (Some(n), 0.0)),
+        }
+    }
+}
+
+pub fn takaheshi_matsuyama<G, P>(graph: &G, nodes: &[G::Node]) -> Option<Vec<Vec<G::Node>>>
+where
+    P: PriorityQueue<QueueObject<Option<G::Node>>> + Debug,
+    P::Handle: Debug,
+    G::Node: Eq + Hash + Clone + Debug,
+    G: WheightedGraph,
+{
+    if nodes.is_empty() {
+        return Some(Vec::new());
+    }
+
+    let mut end_nodes: HashSet<G::Node> = HashSet::from_iter(nodes.iter().cloned());
+
+    let mut start_nodes = vec![end_nodes.iter().next().cloned().unwrap()];
+
+    end_nodes.remove(&start_nodes[0]);
+
+    let mut paths = Vec::new();
+
+    while !end_nodes.is_empty() {
+        let wrapper = MultistartWrapper {
+            graph,
+            start_nodes: &start_nodes,
+        };
+
+        if let Some(p) = dijkstra::<_, P, _>(&wrapper, None, |n| {
+            n.as_ref().is_some_and(|n| end_nodes.contains(n))
+        }) {
+            let p = p.into_iter().skip(1).map(|n| n.unwrap()).collect();
+
+            for n in &p {
+                end_nodes.remove(n);
+                start_nodes.push(n.clone());
+            }
+
+            paths.push(p);
+        } else {
+            return None;
+        }
+    }
+
+    Some(paths)
 }