#864153 (Rust) No.5007 Steiner Space Travel

• 問題: No.5007 Steiner Space Travel
• ユーザー: soso
• 提出日時: 2023-04-29 01:10:39
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 9 ms / 1,000 ms
• コード長: 7,189 bytes
• コンパイル時間: 3,469 ms
• コンパイル使用メモリ: 157,820 KB
• 実行使用メモリ: 4,372 KB
• スコア: 4,300,620
• 最終ジャッジ日時: 2023-04-29 01:10:46
• 合計ジャッジ時間: 6,145 ms
• ジャッジサーバーID （参考情報）: judge15 / judge13

ファイルパターン	結果
other	AC * 30

ソースコード

use std::{cmp::Ordering, collections::BinaryHeap};

const INF: i32 = 1000000000;
const A: i32 = 5;

/// 作問者terry_u16さんの提出から頂いた
/// 入力受け取り用のマクロ
macro_rules! get {
      ($t:ty) => {
          {
              let mut line: String = String::new();
              std::io::stdin().read_line(&mut line).unwrap();
              line.trim().parse::<$t>().unwrap()
          }
      };
      ($($t:ty),*) => {
          {
              let mut line: String = String::new();
              std::io::stdin().read_line(&mut line).unwrap();
              let mut iter = line.split_whitespace();
              (
                  $(iter.next().unwrap().parse::<$t>().unwrap(),)*
              )
          }
      };
      ($t:ty; $n:expr) => {
          (0..$n).map(|_|
              get!($t)
          ).collect::<Vec<_>>()
      };
      ($($t:ty),*; $n:expr) => {
          (0..$n).map(|_|
              get!($($t),*)
          ).collect::<Vec<_>>()
      };
      ($t:ty ;;) => {
          {
              let mut line: String = String::new();
              std::io::stdin().read_line(&mut line).unwrap();
              line.split_whitespace()
                  .map(|t| t.parse::<$t>().unwrap())
                  .collect::<Vec<_>>()
          }
      };
      ($t:ty ;; $n:expr) => {
          (0..$n).map(|_| get!($t ;;)).collect::<Vec<_>>()
      };
}

#[allow(dead_code)]
#[derive(Debug, Clone)]
struct Input {
    n: usize,
    m: usize,
    points: Vec<Point>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum PointType {
    PLANET,
    STATION,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct Point {
    x: i32,
    y: i32,
    point_type: PointType,
}

impl Point {
    fn new(x: i32, y: i32, point_type: PointType) -> Self {
        Self { x, y, point_type }
    }
    fn dist_sq(&self, other: &Self) -> i32 {
        let dx = self.x - other.x;
        let dy = self.y - other.y;
        (dx * dx + dy * dy) as i32
    }
}

fn main() {
    let input = read_input();
    let points = points_add_station(&input.points);

    let distances = warshall_floyd(&points);
    let mut graph = vec![vec![]; input.n];

    for i in 0..input.n {
        for j in 0..input.n {
            if i == j {
                continue;
            }
            graph[i].push(Edge {
                node: j,
                cost: distances[i][j],
            });
            graph[j].push(Edge {
                node: i,
                cost: distances[j][i],
            })
        }
    }

    // 惑星1から出発し、一番近い惑星を貪欲に選び続ける（Nearest Neighbour法）
    let mut v = 0;
    let mut visited = vec![false; input.n];
    visited[0] = true;
    let mut route = vec![0];

    // 惑星1以外のN-1個の惑星を訪問していく
    for _ in 0..input.n - 1 {
        let mut nearest_dist = INF;
        let mut nearest_v = None;

        // 一番近い惑星を探す
        for next in 0..input.n {
            if visited[next] {
                continue;
            }

            if distances[v][next] < nearest_dist {
                nearest_dist = distances[v][next];
                nearest_v = Some(next);
            }
        }

        // パスを復元
        let mut path = dijkstra(&graph, v, nearest_v.unwrap());
        route.append(&mut path);

        // 次の頂点に移動
        v = nearest_v.unwrap();
        visited[v] = true;
    }

    // 最後に惑星1に戻る必要がある
    let mut path = dijkstra(&graph, v, 0);
    route.append(&mut path);

    // 解の出力

    // 宇宙ステーションの座標を出力
    for point in points.iter().skip(input.n) {
        println!("{} {}", point.x, point.y);
    }

    // 経路の長さを出力
    println!("{}", route.len());

    // 経路を出力
    for v in route {
        if v < input.n {
            println!("1 {}", v + 1);
        } else {
            println!("2 {}", v - input.n + 1);
        }
    }
}

/// 入力読み込み
fn read_input() -> Input {
    let (n, m) = get!(usize, usize);

    let mut points = vec![];

    for _ in 0..n {
        let (x, y) = get!(i32, i32);
        points.push(Point::new(x, y, PointType::PLANET));
    }

    Input { n, m, points }
}

fn points_add_station(points: &Vec<Point>) -> Vec<Point> {
    // TODO: ステーションの配置
    let mut new_points = points.clone();

    for i in 0..8 {
        new_points.push(Point {
            x: i,
            y: i,
            point_type: PointType::STATION,
        })
    }

    return new_points;
}

/// エネルギー計算
fn calc_energy(points: &Vec<Point>, i: usize, j: usize) -> i32 {
    let x = points[i];
    let y = points[j];
    let mut energy = x.dist_sq(&y);
    if x.point_type == PointType::PLANET {
        energy *= A;
    }
    if y.point_type == PointType::PLANET {
        energy *= A;
    }

    energy
}

/// ワーシャルフロイド法によって、O(v^3)で全頂点対の最短経路長を求める
fn warshall_floyd(points: &Vec<Point>) -> Vec<Vec<i32>> {
    let mut distances = vec![vec![0; points.len()]; points.len()];
    // 全点間エネルギーを計算
    for i in 0..points.len() {
        for j in 0..points.len() {
            distances[i][j] = calc_energy(points, i, j);
        }
    }
    // ワーシャルフロイド
    for k in 0..points.len() {
        for i in 0..points.len() {
            for j in 0..points.len() {
                let d = distances[i][k] + distances[k][j];
                distances[i][j] = distances[i][j].min(d);
            }
        }
    }
    distances
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
struct State {
    cost: i32,
    position: usize,
}

impl Ord for State {
    fn cmp(&self, other: &Self) -> Ordering {
        other
            .cost
            .cmp(&self.cost)
            .then_with(|| self.position.cmp(&other.position))
    }
}

impl PartialOrd for State {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct Edge {
    node: usize,
    cost: i32,
}

fn dijkstra(graph: &Vec<Vec<Edge>>, start: usize, goal: usize) -> Vec<usize> {
    let mut dist: Vec<_> = (0..graph.len()).map(|_| INF).collect();
    let mut heap = BinaryHeap::new();
    let mut prev_points = vec![None; graph.len()];

    dist[start] = 0;
    heap.push(State {
        cost: 0,
        position: start,
    });

    while let Some(State { cost, position }) = heap.pop() {
        if cost > dist[position] {
            continue;
        }

        for edge in &graph[position] {
            let next = State {
                cost: cost + edge.cost,
                position: edge.node,
            };

            if next.cost < dist[next.position] {
                prev_points[next.position] = Some(position);
                heap.push(next);
                dist[next.position] = next.cost;
            }
        }
    }

    let mut v = goal;
    let mut path = vec![];

    while v != start {
        path.push(v);
        v = prev_points[v].unwrap();
    }

    path.reverse();

    path
}