1 files changed, 25 insertions, 13 deletions

diff --git a/src/tsp_approx.rs b/src/tsp_approx.rs
index 7c9568d..097fee7 100644
--- a/src/tsp_approx.rs
+++ b/src/tsp_approx.rs
@@ -4,7 +4,7 @@ use mset::MultiSet;
 use partitions::partition_vec::PartitionVec;
 use std::{cmp::Eq, hash::Hash, mem::swap};
 
-use crate::{MetricElem, TspAlg};
+use crate::{MetricElem, TspBaseAlg};
 
 struct DistCache(HashMap<(usize, usize), f64>);
 impl DistCache {
@@ -285,9 +285,9 @@ fn christofides(
     r
 }
 
-fn refine_2_opt(dist_cache: &DistCache, tour: Vec<usize>) -> Vec<usize> {
+fn refine_2_opt(dist_cache: &DistCache, tour: Vec<usize>) -> (bool, Vec<usize>) {
     if tour.len() < 4 {
-        return tour;
+        return (false, tour);
     }
 
     // convert the tour into a doubly linked-list. instead of pointers, we use
@@ -327,6 +327,8 @@ fn refine_2_opt(dist_cache: &DistCache, tour: Vec<usize>) -> Vec<usize> {
         }
     };
 
+    let mut improved = false;
+
     // for each combination of edges ...
     while le.0 != n - 2 {
         let mut re = adv(le, &tour);
@@ -355,6 +357,8 @@ fn refine_2_opt(dist_cache: &DistCache, tour: Vec<usize>) -> Vec<usize> {
 
             swap(&mut le.1, &mut re.0);
 
+            improved = true;
+
             re = adv(re, &tour);
         }
 
@@ -372,7 +376,7 @@ fn refine_2_opt(dist_cache: &DistCache, tour: Vec<usize>) -> Vec<usize> {
         cur = adv(cur, &tour);
     }
 
-    tour_flat
+    (improved, tour_flat)
 }
 
 fn get_dist_cache<M>(embeds: &[M], hash_seed: &Option<u64>) -> DistCache
@@ -390,7 +394,12 @@ where
     DistCache(r)
 }
 
-pub(crate) fn tsp<M>(embeds: &[M], alg: &TspAlg, hash_seed: &Option<u64>) -> (Vec<usize>, f64)
+pub(crate) fn tsp<M>(
+    embeds: &[M],
+    alg: &TspBaseAlg,
+    refinements: usize,
+    hash_seed: &Option<u64>,
+) -> (Vec<usize>, f64)
 where
     M: MetricElem,
 {
@@ -405,16 +414,19 @@ where
     let mst = get_mst(&dc, embeds.len(), hash_seed);
 
     bar.set_message("Finding path...");
-    let tour = match alg {
-        TspAlg::MstDfs => tsp_from_mst(mst),
-        TspAlg::Christofides => christofides(&dc, mst, hash_seed),
-        TspAlg::ChristofidesRefined => {
-            let tour = christofides(&dc, mst, hash_seed);
-            bar.set_message("Refining path...");
-            refine_2_opt(&dc, tour)
-        }
+    let mut tour = match alg {
+        TspBaseAlg::MstDfs => tsp_from_mst(mst),
+        TspBaseAlg::Christofides => christofides(&dc, mst, hash_seed),
     };
 
+    for _ in 0..refinements {
+        let res = refine_2_opt(&dc, tour);
+        tour = res.1;
+        if !res.0 {
+            break; // stop early at convergence
+        }
+    }
+
     let mut total_dist = 0.;
     for i in 0..tour.len() - 1 {
         let (a, b) = (tour[i], tour[i + 1]);