TSP.rs

tsp.rs

#[macro_use]
extern crate log;
#[macro_use]
extern crate serde_derive;
extern crate clap;
extern crate csv;
extern crate loggerv;
extern crate ordered_float;
extern crate rand;
extern crate rayon;
extern crate serde;

use clap::{App, Arg};
use ordered_float::OrderedFloat;
use rand::{Rng, SeedableRng, StdRng};
use rayon::prelude::*;
use std::cmp;
use std::collections::HashSet;
use std::path::Path;

type Goal = Option<f64>;
type DistanceMatrix = Vec<Vec<f64>>;
type Route = Vec<usize>;

#[derive(Deserialize, Clone)]
struct City {
    x: f64,
    y: f64,
}

enum Strategy {
    Greedy,
    DP,
    SA,
    TwoOpt(Goal),
    MapReduce,
}

struct TSP {
    cities: Vec<City>,
    strategy: Strategy,
}

fn build_cli() -> App<'static, 'static> {
    App::new("tsp")
        .version("1.0")
        .arg(Arg::with_name("v").short("v").multiple(true).help(
            "Sets the level of verbosity",
        ))
        .arg(
            Arg::with_name("strategy")
                .short("s")
                .long("strategy")
                .takes_value(true)
                .required(true)
                .possible_values(&["greedy", "dp", "sa", "2opt", "speed", "mapreduce"]),
        )
        .arg(Arg::with_name("input").required(true))
}

fn main() {
    let matches = build_cli().get_matches();
    loggerv::init_with_verbosity(matches.occurrences_of("v")).unwrap();
    let cities = TSP::read_from_csv(matches.value_of("input").unwrap());
    let strategy = matches.value_of("strategy").unwrap();

    let tsp = TSP::new(
        cities,
        match strategy {
            "greedy" => Strategy::Greedy,
            "dp" => Strategy::DP,
            "sa" => Strategy::SA,
            "2opt" => Strategy::TwoOpt(None),
            "speed" => Strategy::TwoOpt(Some(47_000.0)),
            "mapreduce" => Strategy::MapReduce,
            _ => unreachable!(),
        },
    );
    let route = tsp.solve();
    debug_assert_eq!(
        route.iter().cloned().collect::<HashSet<usize>>().len(),
        tsp.cities.len()
    );
    println!("index");
    for i in &route {
        println!("{}", i);
    }
    info!("Distance: {}", tsp.route_dist(&route));
}

impl TSP {
    fn read_from_csv<P: AsRef<Path>>(input: P) -> Vec<City> {
        let mut reader = csv::Reader::from_path(input).unwrap();
        reader.deserialize().map(|r| r.unwrap()).collect()
    }

    fn new(cities: Vec<City>, strategy: Strategy) -> TSP {
        TSP { cities, strategy }
    }

    fn distance_matrix(&self) -> DistanceMatrix {
        let n = self.cities.len();
        let mut matrix = vec![vec![0.0; n]; n];
        for i in 0..n {
            for j in 0..i {
                matrix[i][j] = self.dist(i, j);
                matrix[j][i] = matrix[i][j];
            }
        }
        matrix
    }

    fn dist(&self, i: usize, j: usize) -> f64 {
        (self.cities[i].x - self.cities[j].x).hypot(self.cities[i].y - self.cities[j].y)
    }

    fn route_dist(&self, route: &Route) -> f64 {
        route.iter().fold(0.0, |acc, &i| {
            acc + self.dist(route[i], route[(i + 1) % self.cities.len()])
        })
    }

    fn solve(&self) -> Route {
        match self.strategy {
            Strategy::Greedy => self.greedy(),
            Strategy::DP => self.dp(),
            Strategy::SA => self.sa(),
            Strategy::TwoOpt(goal) => self.two_opt(goal),
            Strategy::MapReduce => self.map_reduce(),
        }
    }

    fn greedy(&self) -> Route {
        let mut unvisited = (1..self.cities.len()).collect::<HashSet<usize>>();
        let mut current = 0;
        let mut route = vec![current];
        while let Some(&i) = unvisited.iter().min_by_key(
            |&&i| OrderedFloat(self.dist(current, i)),
        )
        {
            route.push(i);
            unvisited.remove(&i);
            current = i;
        }
        route
    }

    fn dp(&self) -> Route {
        let dist = self.distance_matrix();
        let n = self.cities.len();
        let mut dp = vec![vec![std::f64::MAX; 1 << n]; n];
        let mut backtrack = vec![vec![0; 1 << n]; n];

        for i in 0..n {
            dp[i][0] = dist[i][0];
        }

        for unvisited in 0..(1 << n) {
            for current in 0..n {
                if (unvisited & (1 << current)) != 0 {
                    continue;
                }
                for next in 0..n {
                    if (unvisited & (1 << next)) == 0 {
                        continue;
                    }
                    let next_unvisited = unvisited ^ (1 << next);
                    let score = dp[next][next_unvisited] + dist[current][next];
                    if score < dp[current][unvisited] {
                        dp[current][unvisited] = score;
                        backtrack[current][unvisited] = next;
                    }
                }
            }
        }
        let mut current = 0;
        let mut unvisited = (1 << n) - 1;
        unvisited ^= 1 << current;
        debug!("Distance (dp): {}", dp[current][unvisited]);
        let mut route = vec![current];
        while unvisited != 0 {
            current = backtrack[current][unvisited];
            route.push(current);
            unvisited ^= 1 << current;
        }
        route
    }

    fn sa(&self) -> Route {

        fn next_route<R: Rng>(n: usize, dist: &DistanceMatrix, route: &Route, rng: &mut R) -> (f64, usize, usize) {
            let i = rng.gen_range(0, n);
            let j = rng.gen_range(0, n);
            let (i, j) = (cmp::min(i, j), cmp::max(i, j));
            if i + 1 < j {
                let i2 = if i == 0 { n - 1 } else { i - 1 };
                let j2 = j - 1;
                let dist_add = dist[route[i2]][route[j2]] + dist[route[i]][route[j]];
                let dist_sub = dist[route[i2]][route[i]] + dist[route[j2]][route[j]];
                let dist_diff = dist_add - dist_sub;
                (dist_diff, i, j)
            } else {
                next_route(n, dist, route, rng)
            }
        }

        let dist = self.distance_matrix();
        let seed: &[_] = &[1, 2, 3, 4];
        let mut rng: StdRng = SeedableRng::from_seed(seed);

        let n = self.cities.len();
        let mut route = (0..n).collect::<Route>();
        let mut iteration = 0;
        let mut temperature = 100000.0;
        let cooling_rate = 0.9999999;
        let absolute_temperature = 0.000001;

        while temperature > absolute_temperature {
            if iteration % 1_000_000 == 0 {
                debug!(
                    "iteration: {}, temperature: {:10.8}, route_dist: {}",
                    iteration,
                    temperature,
                    self.route_dist(&route)
                );
            }
            let (delta, i, j) = next_route(n, &dist, &route, &mut rng);
            if delta < 0.0 || (-delta / temperature).exp() > rng.next_f64() {
                route[i..j].reverse();
            }
            temperature *= cooling_rate;
            iteration += 1;
        }
        route
    }

    fn two_opt(&self, goal: Option<f64>) -> Route {
        let dist = self.distance_matrix();
        let mut route = self.greedy();
        let n = self.cities.len();
        let mut total_dist = self.route_dist(&route);
        loop {
            let mut improved = false;
            for i in 0..n {
                for j in (i + 2)..n {
                    let i2 = if i == 0 { n - 1 } else { i - 1 };
                    let j2 = j - 1;
                    let dist_add = dist[route[i2]][route[j2]] + dist[route[i]][route[j]];
                    let dist_sub = dist[route[i2]][route[i]] + dist[route[j2]][route[j]];
                    let dist_diff = dist_add - dist_sub;
                    if dist_diff < 0.0 {
                        improved = true;
                        route[i..j].reverse();
                        if let Some(goal) = goal {
                            total_dist += dist_diff;
                            if total_dist < goal {
                                return route;
                            }
                        }
                    }
                }
            }
            if !improved {
                return route;
            }
        }
    }

    fn map_reduce(&self) -> Route {
        #[derive(Clone)]
        struct Block {
            cities: Vec<City>,
            city_indices: Vec<usize>,
        }

        impl Block {
            fn new() -> Block {
                Block {
                    cities: Vec::new(),
                    city_indices: Vec::new(),
                }
            }
        }

        // TODO(hayato): Avoid magic numbers
        let dx = 1600;
        let dy = 900;
        let block_size = 25;
        assert_eq!(dx % block_size, 0);
        assert_eq!(dy % block_size, 0);
        let bx = dx / block_size;
        let by = dy / block_size;
        let blocks_num = bx * by;
        let mut blocks = vec![Block::new(); blocks_num];
        for (i, city) in self.cities.iter().enumerate() {
            let x = (city.x as usize) / block_size;
            let y = (city.y as usize) / block_size;
            let block_index = y * bx + x;
            blocks[block_index].city_indices.push(i);
            blocks[block_index].cities.push(city.clone());
        }

        let mut routes = blocks
            .into_par_iter()
            .map(|Block {
                 cities,
                 city_indices,
             }| {
                TSP::new(cities, Strategy::Greedy)
                    .solve()
                    .into_iter()
                    .map(|i| city_indices[i])
                    .collect::<Route>()
            })
            .collect::<Vec<_>>();

        let mut route = Vec::with_capacity(self.cities.len());

        // TODO(hayato): Use hilbert curve
        assert_eq!(bx % 2, 0);
        assert_eq!(by % 2, 0);
        let hx = bx / 2;
        for y in 0..by {
            if y % 2 == 0 {
                for x in (0..hx).rev() {
                    route.append(&mut routes[y * bx + x]);
                }
            } else {
                for x in 0..hx {
                    route.append(&mut routes[y * bx + x]);
                }
            }
        }
        for y in (0..by).rev() {
            if y % 2 == 0 {
                for x in (hx..bx).rev() {
                    route.append(&mut routes[y * bx + x]);
                }
            } else {
                for x in hx..bx {
                    route.append(&mut routes[y * bx + x]);
                }
            }
        }
        route
    }
}