steiner-tree 0.0.1

// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: GPL-3.0-or-later

//! Approximate solver for the travelling salesperson problem (aka. TSP, travelling salesman problem).

// /// Define a distance metric.
// pub trait DistanceMetric {
//     type Distance: Ord;
//
//     /// Measure the distance between two elements.
//     fn distance(&self, other: &Self) -> Self::Distance;
// }

use num_traits::{Num, Zero};
use std::ops::{Add, Sub};

/// Represent the tour using an immutable array of points,
/// a `route` array of indices into the points
/// and a distance function.
pub struct Tour<'a, T, DFn> {
    distance: &'a DFn,
    points: &'a [T],
    route: Vec<usize>,
}

impl<'a, T, DFn, D> Tour<'a, T, DFn>
    where DFn: Fn(&T, &T) -> D,
          D: Ord {
    fn new(points: &'a [T], distance: &'a DFn) -> Self {
        Self {
            distance,
            points,
            route: (0..points.len()).collect(),
        }
    }

    fn new_with_route(points: &'a [T], route: Vec<usize>, distance: &'a DFn) -> Self {
        assert_eq!(points.len(), route.len());
        Self {
            distance,
            points,
            route,
        }
    }

    /// Number of points.
    fn len(&self) -> usize { self.points.len() }

    /// Get the index of the next element in the tour (wraps around).
    fn next_index(&self, idx: usize) -> usize {
        assert!(idx < self.len());
        if idx + 1 < self.len() {
            idx + 1
        } else {
            0
        }
    }

    fn prev_index(&self, idx: usize) -> usize {
        assert!(idx < self.len());
        if idx == 0 {
            self.len() - 1
        } else {
            idx - 1
        }
    }

    /// Get a reference to the element which is currently at the `index` in the tour.
    fn element(&self, index: usize) -> &T {
        &self.points[self.route[index]]
    }

    fn elements(&self, idx1: usize, idx2: usize) -> (&T, &T) {
        (self.element(idx1), self.element(idx2))
    }

    /// Get elements at start and end of the edge.
    fn edge_elements(&self, edge_start_idx: usize) -> (&T, &T) {
        self.elements(edge_start_idx, self.next_index(edge_start_idx))
    }

    /// Get the length of the edge. Edge indices start with zero.
    /// The last edge wraps around to the first element of the tour.
    fn edge_length(&self, edge_start_idx: usize) -> D {
        let (a, b) = self.edge_elements(edge_start_idx);
        (self.distance)(a, b)
    }

}

impl<'a, T, DFn, D> Tour<'a, T, DFn>
    where DFn: Fn(&T, &T) -> D,
          D: Ord + Add<Output=D> + Zero {

    /// Compute total length of tour without counting the edge from the last point to the first.
    fn total_length_nowrap(&self) -> D {
        (1..self.len())
            .map(|i| self.edge_length(i-1))
            .fold(Zero::zero(), |l, acc| l + acc)
    }


    /// Get sum of the edges adjacent to the vertex at `idx`.
    /// The edge which wraps from the last element to the first is ignored.
    fn length_of_adjacent_edges_nowrap(&self, idx: usize) -> D {
        if idx == 0 {
            self.edge_length(idx)
        } else if idx == self.len() - 1 {
            self.edge_length(self.prev_index(idx))
        } else {
            self.edge_length(idx) + self.edge_length(self.prev_index(idx))
        }
    }

}

pub fn solve_tsp<T, DFn, D>(points: &[T], distance_fn: &DFn) -> Vec<usize>
    where DFn: Fn(&T, &T) -> D,
          D: Ord + Add<Output=D> + Sub<Output=D> + Zero {

    // Find initial tour by greedy nearest-neighbour search.
    let route = solve_tsp_nearest_neighbour(points, distance_fn);
    let mut tour = Tour::new_with_route(points, route, distance_fn);
    // let mut tour = Tour::new(points, distance_fn);
    let initial_cost = tour.total_length_nowrap();

    // Check if the total length gets improved when swapping the both points.
    let swap_improves_length = |tour: &Tour<T, DFn>, idx1: usize, idx2: usize| -> bool {
        debug_assert!(idx1 < tour.len());
        debug_assert!(idx2 < tour.len());

        if idx1 == idx2 {
            return false;
        }

        // Sort indices.
        let (idx1, idx2) = if idx1 <= idx2 {
            (idx1, idx2)
        } else {
            (idx2, idx1)
        };

        debug_assert_ne!(idx1 + 1, tour.len());
        debug_assert_ne!(idx2, 0);

        let p1 = tour.element(idx1);
        let p2 = tour.element(idx2);
        let p1_prev = tour.element(tour.prev_index(idx1));
        let p2_prev = tour.element(tour.prev_index(idx2));
        let p1_next = tour.element(tour.next_index(idx1));
        let p2_next = tour.element(tour.next_index(idx2));


        let d_now1 = if idx1 == 0 { Zero::zero() } else { distance_fn(p1_prev, p1) }
            + if idx1 + 1 == tour.len() || idx1 + 1 == idx2 { Zero::zero() } else { distance_fn(p1, p1_next) };

        let d_now2 = if idx1 + 1 == idx2 { Zero::zero() } else { distance_fn(p2_prev, p2) }
            + if idx2 + 1 == tour.len() { Zero::zero() } else { distance_fn(p2, p2_next) };

        let d_swapped1 = if idx1 == 0 { Zero::zero() } else { distance_fn(p1_prev, p2) }
            + if idx1 + 1 == tour.len() || idx1 + 1 == idx2 { Zero::zero() } else { distance_fn(p2, p1_next) };

        let d_swapped2 = if idx1 + 1 == idx2 { Zero::zero() } else { distance_fn(p2_prev, p1) }
            + if idx2 + 1 == tour.len() { Zero::zero() } else { distance_fn(p1, p2_next) };

        let d_now = d_now1 + d_now2;
        let d_swapped = d_swapped1 + d_swapped2;

        d_now > d_swapped
    };

    // Try to improve the length by swapping pairs of points.
    'outer: loop {
        'inner: for i in 0..tour.len() {
            for j in i..tour.len() {

                if swap_improves_length(&tour, i, j) {
                    // Swap i and j.
                    tour.route.swap(i, j);
                    break 'inner;
                }
            }
            if i == tour.len() - 1 {
                // Reached end without finding any improvement.
                break 'outer;
            }
        }
    }

    // let mut cant_improve = vec![false; points.len()];
    // for i in 0..10 {
    //
    //     let swap_candidate1 = (1..points.len()-1)
    //         .filter(|i| !cant_improve[*i])
    //         .max_by_key(|i| tour.length_of_adjacent_edges_nowrap(*i))
    //         .unwrap();
    //     // Find best improvement.
    //     let swap_candidate2 = (1..points.len()-1)
    //         .find(|i| swap_improves_length(&tour, swap_candidate1, *i));
    //     if let Some(swap_candidate2) = swap_candidate2 {
    //         cant_improve[swap_candidate1] = false;
    //         cant_improve[swap_candidate2] = false;
    //         cant_improve[swap_candidate2+1] = false;
    //         cant_improve[swap_candidate2-1] = false;
    //         tour.route.swap(swap_candidate1, swap_candidate2);
    //     } else {
    //         cant_improve[swap_candidate1] = true;
    //     }
    // }

    let cost = tour.total_length_nowrap();

    debug_assert!(cost <= initial_cost);

    tour.route
}

/// Find approximate shortest tour through the `points` by a greedy
/// iterative shortest-neighbour search.
fn solve_tsp_nearest_neighbour<T, DFn, D>(points: &[T], distance_fn: &DFn) -> Vec<usize>
    where DFn: Fn(&T, &T) -> D,
          D: Ord {
    let mut tour = Tour::new(points, distance_fn);

    if tour.len() <= 2 {
        return tour.route;
    }

    // Greedy nearest-neighbour search.
    for i in 0..tour.len() - 1 {
        // Find nearest neighbour of the current node.
        let nearest_neighbour_idx = tour.route[i + 1..].iter()
            .copied()
            .min_by_key(|&j| tour.edge_length(j))
            .unwrap();

        // Move nearest neighbour to position i+1.
        tour.route.swap(i + 1, nearest_neighbour_idx);
    }

    // Break the tour at the longest edge.
    let start_index_of_longest_edge = (0..tour.len())
        .max_by_key(|&a| tour.edge_length(a))
        .unwrap();
    let end_index_of_longest_edge = tour.next_index(start_index_of_longest_edge);
    tour.route.rotate_left(end_index_of_longest_edge);

    tour.route
}

#[test]
fn test_solve_tsp_nearest_neighbour() {
    let points = vec![(2i32, 2i32), (0, 0), (1, 1)];

    let tour = solve_tsp_nearest_neighbour(
        &points,
        &|(ax, ay), (bx, by)| (ax - bx).abs() + (ay - by).abs(),
    );

    assert_eq!(tour, [1, 2, 0]);
}

#[test]
fn test_solve_tsp() {
    let points = vec![(2i32, 2i32), (0, 0), (1, 1)];

    let tour = solve_tsp(
        &points,
        &|(ax, ay), (bx, by)| (ax - bx).abs() + (ay - by).abs(),
    );

    assert_eq!(tour, [1, 2, 0]);
}