Use connected set cover for roboports

factorio-blueprint/src/abstraction/roboports.rs

@@ -1,12 +1,32 @@
 use factorio_core::{misc::PositionMap, prelude::Position};
 use factorio_graph::{
-    priority_queue::binary_heap::FastBinaryHeap, set_cover::greedy_set_cover_priority_queue,
+    graph::Graph,
+    priority_queue::binary_heap::FastBinaryHeap,
+    set_cover::{greedy_connected_set_cover_priority_queue, greedy_set_cover_priority_queue},
 };
 
 use crate::abstraction::Entity;
 
 use super::Blueprint;
 
+struct RoboportGraph<'a> {
+    roboports: &'a [Position],
+}
+
+impl Graph for RoboportGraph<'_> {
+    type Node = usize;
+
+    fn edge(&self, node: &Self::Node, num: usize) -> Option<Self::Node> {
+        let p = self.roboports[*node];
+        self.roboports
+            .iter()
+            .enumerate()
+            .filter(|(_, other)| p.x.abs_diff(other.x) <= 100 && p.y.abs_diff(other.y) <= 100)
+            .nth(num)
+            .map(|(i, _)| i)
+    }
+}
+
 impl Blueprint {
     pub fn add_roboport_network(&mut self) {
         let aabb = self.get_aabb().unwrap();
@@ -41,7 +61,13 @@ impl Blueprint {
         }
 
         dbg!(universe, sets.len());
-        let res = greedy_set_cover_priority_queue::<_, FastBinaryHeap<_>>(universe, &sets);
+        let res = greedy_connected_set_cover_priority_queue::<_, _, FastBinaryHeap<_>>(
+            universe,
+            &sets,
+            RoboportGraph {
+                roboports: &roboports,
+            },
+        );
 
         for r in res {
             self.add_entity(Entity::new_roboport(roboports[r]));

factorio-graph/src/graph.rs

@@ -0,0 +1,47 @@
+pub trait Graph: Sized {
+    type Node;
+
+    fn num_edges(&self, node: &Self::Node) -> usize {
+        self.edge_iter(node).count()
+    }
+
+    fn edge(&self, node: &Self::Node, num: usize) -> Option<Self::Node>;
+
+    fn edge_iter<'a, 'b>(&'a self, node: &'b Self::Node) -> GraphEdgeIter<'a, 'b, Self> {
+        GraphEdgeIter::new(self, node)
+    }
+}
+
+pub struct GraphEdgeIter<'a, 'b, G>
+where
+    G: Graph,
+{
+    graph: &'a G,
+    node: &'b G::Node,
+    count: usize,
+}
+
+impl<'a, 'b, G> GraphEdgeIter<'a, 'b, G>
+where
+    G: Graph,
+{
+    pub fn new(graph: &'a G, node: &'b G::Node) -> Self {
+        Self {
+            graph,
+            node,
+            count: 0,
+        }
+    }
+}
+
+impl<G> Iterator for GraphEdgeIter<'_, '_, G>
+where
+    G: Graph,
+{
+    type Item = G::Node;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.count += 1;
+        self.graph.edge(self.node, self.count - 1)
+    }
+}

factorio-graph/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod graph;
 pub mod priority_queue;
 pub mod set_cover;
 pub mod wheighted_graph;

factorio-graph/src/priority_queue/binary_heap.rs

@@ -91,6 +91,13 @@ where
         }
     }
 
+    fn increase_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item)) {
+        if let Some(index) = self.search(*handle) {
+            f(&mut self.data[index].item);
+            self.downheap(index);
+        }
+    }
+
     fn new() -> Self {
         Self {
             data: Vec::new(),
@@ -134,6 +141,14 @@ where
             println!("Decrease key: {old_i:?} -> {i:?}");
         })
     }
+
+    fn increase_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut I)) {
+        self.inner.increase_key(handle, |i| {
+            let old_i = i.clone();
+            f(i);
+            println!("Increase key: {old_i:?} -> {i:?}");
+        })
+    }
 }
 
 #[derive(Debug)]
@@ -221,6 +236,13 @@ where
         }
     }
 
+    fn increase_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item)) {
+        if let Some(index) = self.search(*handle) {
+            f(&mut self.data[index].item);
+            self.downheap(index);
+        }
+    }
+
     fn new() -> Self {
         Self {
             data: Vec::new(),

factorio-graph/src/priority_queue/fibonacci_heap.rs

@@ -275,4 +275,10 @@ where
         f(&mut h.item);
         self.cut(handle.0);
     }
+
+    fn increase_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut I)) {
+        let h = self.get_mut(handle.0);
+        f(&mut h.item);
+        self.cut(handle.0);
+    }
 }

factorio-graph/src/priority_queue/mod.rs

@@ -14,6 +14,7 @@ where
     fn insert(&mut self, item: Item) -> Self::Handle;
     fn pop_min(&mut self) -> Option<Item>;
     fn decrease_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item));
+    fn increase_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item));
 }
 
 #[cfg(test)]

factorio-graph/src/set_cover.rs

@@ -1,6 +1,6 @@
 use std::ops::Deref;
 
-use crate::priority_queue::PriorityQueue;
+use crate::{graph::Graph, priority_queue::PriorityQueue};
 
 #[derive(Debug, Clone, Copy)]
 pub struct SetUncovered {
@@ -100,7 +100,7 @@ where
             }
 
             // dbg!(decrease, i, s.deref());
-            p.decrease_key(h, |e| e.uncoveredElements -= decrease);
+            p.increase_key(h, |e| e.uncoveredElements -= decrease);
         }
 
         for &e in sets[i].deref() {
@@ -114,6 +114,103 @@ where
     r
 }
 
+enum HandleOption<H> {
+    Unknown,
+    Seen(H),
+    Visited,
+}
+
+pub fn greedy_connected_set_cover_priority_queue<S, G, P>(
+    universe: usize,
+    sets: &[S],
+    graph: G,
+) -> Vec<usize>
+where
+    S: Deref<Target = [usize]>,
+    G: Graph<Node = usize>,
+    P: PriorityQueue<SetUncovered>,
+{
+    let mut covered = vec![false; universe];
+    let mut covered_count = 0;
+
+    let mut p = P::new();
+
+    let mut handles = (0..sets.len())
+        .map(|_| HandleOption::Unknown)
+        .collect::<Vec<_>>();
+
+    let mut r = Vec::new();
+
+    let (i, _s) = sets
+        .iter()
+        .enumerate()
+        .max_by_key(|(_, s)| s.len())
+        .unwrap();
+
+    r.push(i);
+    handles[i] = HandleOption::Visited;
+
+    for &e in sets[i].deref() {
+        if !covered[e] {
+            covered_count += 1;
+            covered[e] = true;
+        }
+    }
+    for e in graph.edge_iter(&i) {
+        handles[e] = HandleOption::Seen(p.insert(SetUncovered {
+            setIndex: e,
+            uncoveredElements: sets[e].iter().filter(|&&v| !covered[v]).count(),
+        }));
+    }
+
+    while covered_count < universe {
+        let SetUncovered {
+            setIndex: i,
+            uncoveredElements: _,
+        } = p.pop_min().unwrap();
+
+        r.push(i);
+        handles[i] = HandleOption::Visited;
+
+        for (h, s) in handles
+            .iter()
+            .zip(sets.iter())
+            .filter_map(|(h, s)| match h {
+                HandleOption::Unknown => None,
+                HandleOption::Seen(h) => Some((h, s)),
+                HandleOption::Visited => None,
+            })
+        {
+            let mut decrease = 0;
+            for &e in sets[i].deref() {
+                if !covered[e] && s.contains(&e) {
+                    decrease += 1;
+                }
+            }
+
+            // dbg!(decrease, i, s.deref());
+            p.increase_key(h, |e| e.uncoveredElements -= decrease);
+        }
+
+        for &e in sets[i].deref() {
+            if !covered[e] {
+                covered_count += 1;
+                covered[e] = true;
+            }
+        }
+        for e in graph.edge_iter(&i) {
+            if matches!(handles[e], HandleOption::Unknown) {
+                handles[e] = HandleOption::Seen(p.insert(SetUncovered {
+                    setIndex: e,
+                    uncoveredElements: sets[e].iter().filter(|&&v| !covered[v]).count(),
+                }));
+            }
+        }
+    }
+
+    r
+}
+
 #[cfg(test)]
 mod test {
     use crate::{