tsp from mst, remove some debugging, print final path

2 files changed, 26 insertions, 16 deletions

diff --git a/src/embedders.rs b/src/embedders.rs
index a8d34b4..0caf2c8 100644
--- a/src/embedders.rs
+++ b/src/embedders.rs
@@ -24,7 +24,6 @@ impl EmbedderT for BrightnessEmbedder {
         paths
             .iter()
             .map(|p| {
-                println!("{:?}", p);
                 let im = image::open(p).map_err(|e| e.to_string())?;
                 let num_bytes = 3 * (im.height() * im.width());
 
diff --git a/src/main.rs b/src/main.rs
index 5cfc10c..26e2b4e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,7 +3,7 @@
 use clap::Parser;
 use priority_queue::PriorityQueue;
 use std::cmp::Ordering;
-use std::collections::HashSet;
+use std::collections::HashMap;
 use std::path::PathBuf;
 
 use embedders::*;
@@ -36,7 +36,7 @@ struct Args {
 //    Ok(Config{base_dir: dirs})
 //}
 
-fn process_embedder<E>(mut e: E, args: Args) -> Result<(), String>
+fn process_embedder<E>(mut e: E, args: Args) -> Result<Vec<PathBuf>, String>
     where E: EmbedderT
 {
     // wrapper struct to
@@ -57,7 +57,7 @@ fn process_embedder<E>(mut e: E, args: Args) -> Result<(), String>
 
     let num_ims = args.images.len();
     if num_ims == 0 {
-        return Ok(());
+        return Ok(Vec::new());
     }
 
     let embeds: Vec<_> = e
@@ -67,14 +67,13 @@ fn process_embedder<E>(mut e: E, args: Args) -> Result<(), String>
 
     let mut possibleEdges =
         PriorityQueue::with_capacity((num_ims * num_ims - num_ims) / 2);
-    let mut visitedImages = HashSet::with_capacity(num_ims);
-    let mut takenEdges = HashSet::with_capacity(num_ims - 1); // TODO das muss anders
+    let mut visitedImages = HashMap::with_capacity(num_ims);
 
     // here, we start at 0.
     // we might get a better result in the end if we started with a vertex next
     // to the lowest-cost edge, but we don't know which one that is (though we
     // could compute that without changing our asymptotic complexity)
-    visitedImages.insert(0);
+    visitedImages.insert(0, Vec::new());
     for i in 1..num_ims {
         possibleEdges.push((0, i), DownOrd(embeds[0].dist(&embeds[i])));
     }
@@ -84,22 +83,20 @@ fn process_embedder<E>(mut e: E, args: Args) -> Result<(), String>
         // find the edge with the least cost that connects us to a new vertex
         let (new, old) = loop {
             let ((a, b), _) = possibleEdges.pop().unwrap();
-            if !visitedImages.contains(&a) {
+            if !visitedImages.contains_key(&a) {
                 break (a, b);
             }
-            else if !visitedImages.contains(&b) {
+            else if !visitedImages.contains_key(&b) {
                 break (b, a);
             }
         };
-        visitedImages.insert(new);
-
-        println!("next edge: {:?} -> {:?}", old, new);
+        visitedImages.insert(new, Vec::new());
 
         // insert all the new edges we could take
-        takenEdges.insert((new, old));
+        visitedImages.entry(old).and_modify(|v|v.push(new));
         for i in 0..num_ims {
             // don't consider edges taking us to nodes we already visited
-            if visitedImages.contains(&i) {
+            if visitedImages.contains_key(&i) {
                 continue;
             }
 
@@ -107,15 +104,29 @@ fn process_embedder<E>(mut e: E, args: Args) -> Result<(), String>
         }
     }
 
-    Ok(())
+    // find TSP approximation via DFS through the MST
+    fn dfs(cur: usize, t: &HashMap<usize, Vec<usize>>, into: &mut Vec<usize>) {
+        into.push(cur);
+        t.get(&cur).unwrap().iter().for_each(|c| dfs(*c, t, into));
+    }
+    let mut tsp_path = Vec::with_capacity(num_ims);
+    dfs(0, &visitedImages, &mut tsp_path);
+
+    Ok(tsp_path.iter().map(|i| args.images[*i].clone()).collect())
 }
 
 fn main() -> Result<(), String> {
     //let cfg = get_config()?;
     let args = Args::parse();
 
-    match args.embedder {
+    let tsp_path = match args.embedder {
         Embedder::Brightness => process_embedder(BrightnessEmbedder, args),
         Embedder::Hue => process_embedder(HueEmbedder, args),
+    }?;
+
+    for p in tsp_path {
+        println!("{:?}", p);
     }
+
+    Ok(())
 }