#1169513 (Rust) No.5007 Steiner Space Travel

• 問題: No.5007 Steiner Space Travel
• ユーザー: assy1028
• 提出日時: 2026-06-21 10:28:07
• 言語: Rust (1.94.0 + proconio + num + itertools)
• 結果: AC
• 実行時間: 642 ms / 1,000 ms
• コード長: 23,513 bytes
• コンパイル時間: 11,471 ms
• コンパイル使用メモリ: 211,944 KB
• 実行使用メモリ: 6,400 KB
• スコア: 8,925,398
• 最終ジャッジ日時: 2026-06-21 10:28:40
• 合計ジャッジ時間: 25,188 ms
• ジャッジサーバーID （参考情報）: judge1_0 / judge2_0

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

ソースコード

use proconio::input;
use std::time::{Duration, Instant};

const PLANET_WEIGHT: f64 = 5.0;
const INF: f64 = 1.0e100;
const SOFT_LIMIT: Duration = Duration::from_millis(620);

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

#[derive(Clone)]
struct Evaluation {
    cost: f64,
    stations: Vec<Point>,
    route: Vec<usize>,
    next: Vec<Vec<usize>>,
}

struct XorShift {
    state: u64,
}

impl XorShift {
    fn new(seed: u64) -> Self {
        let state = if seed == 0 {
            0x9e37_79b9_7f4a_7c15
        } else {
            seed
        };
        Self { state }
    }

    fn next_u64(&mut self) -> u64 {
        let mut x = self.state;
        x ^= x << 7;
        x ^= x >> 9;
        self.state = x;
        x
    }

    fn gen_usize(&mut self, upper: usize) -> usize {
        if upper <= 1 {
            0
        } else {
            (self.next_u64() as usize) % upper
        }
    }
}

fn main() {
    input! {
        n: usize,
        m: usize,
        coords: [(i32, i32); n],
    }

    let planets: Vec<Point> = coords.into_iter().map(|(x, y)| Point { x, y }).collect();
    let answer = solve(&planets, m);

    let mut out = String::new();
    for p in &answer.stations {
        out.push_str(&format!("{} {}\n", p.x, p.y));
    }
    out.push_str(&format!("{}\n", answer.path.len()));
    for node in answer.path {
        if node < planets.len() {
            out.push_str(&format!("1 {}\n", node + 1));
        } else {
            out.push_str(&format!("2 {}\n", node + 1 - planets.len()));
        }
    }
    print!("{out}");
}

struct Answer {
    stations: Vec<Point>,
    path: Vec<usize>,
}

fn solve(planets: &[Point], station_count: usize) -> Answer {
    let started_at = Instant::now();
    let mut rng = XorShift::new(input_hash(planets) ^ station_count as u64);

    let mut top_evals: Vec<Evaluation> = Vec::new();
    let trial_limit = if planets.len() <= 10 { 24 } else { 72 };
    for trial in 0..trial_limit {
        if started_at.elapsed() >= SOFT_LIMIT {
            break;
        }
        let stations = kmeans_stations(planets, station_count, trial, &mut rng);
        let eval = evaluate(planets, stations);
        push_top_eval(&mut top_evals, eval, 6);
    }

    if top_evals.is_empty() {
        let stations = fallback_stations(planets, station_count);
        push_top_eval(&mut top_evals, evaluate(planets, stations), 1);
    }

    let mut best = top_evals[0].clone();

    for mut cur in top_evals {
        for _ in 0..4 {
            let path = expand_route(planets.len(), &cur.route, &cur.next);
            let stations = optimize_stations_for_path(planets, &cur.stations, &path);
            let eval = evaluate(planets, stations);
            if eval.cost + 1.0e-9 < cur.cost {
                cur = eval;
            } else {
                break;
            }
        }
        if cur.cost < best.cost {
            best = cur;
        }
    }

    if started_at.elapsed() < SOFT_LIMIT {
        let mut stations = best.stations.clone();
        improve_stations_for_route(planets, &mut stations, &best.route, started_at);
        let eval = evaluate(planets, stations);
        if eval.cost < best.cost {
            best = eval;
        }
    }

    let path = expand_route(planets.len(), &best.route, &best.next);
    Answer {
        stations: best.stations,
        path,
    }
}

fn push_top_eval(top: &mut Vec<Evaluation>, eval: Evaluation, keep: usize) {
    top.push(eval);
    top.sort_by(|a, b| a.cost.partial_cmp(&b.cost).unwrap());
    top.truncate(keep);
}

fn input_hash(points: &[Point]) -> u64 {
    let mut h = 0xcbf2_9ce4_8422_2325_u64;
    for p in points {
        h ^= p.x as u64 + 0x9e37_79b9;
        h = h.wrapping_mul(0x1000_0000_01b3);
        h ^= p.y as u64 + 0x85eb_ca6b;
        h = h.wrapping_mul(0x1000_0000_01b3);
    }
    h
}

fn fallback_stations(planets: &[Point], station_count: usize) -> Vec<Point> {
    if planets.is_empty() {
        return vec![Point { x: 500, y: 500 }; station_count];
    }
    let mut stations = Vec::with_capacity(station_count);
    let mut chosen = vec![false; planets.len()];
    stations.push(planets[0]);
    chosen[0] = true;

    while stations.len() < station_count {
        let mut best_i = 0;
        let mut best_d = -1.0;
        for (i, &p) in planets.iter().enumerate() {
            if chosen[i] && stations.len() < planets.len() {
                continue;
            }
            let d = stations
                .iter()
                .map(|&s| sq_dist_i(p, s))
                .fold(INF, f64::min);
            if d > best_d {
                best_d = d;
                best_i = i;
            }
        }
        stations.push(planets[best_i]);
        chosen[best_i] = true;
    }
    stations
}

fn kmeans_stations(
    planets: &[Point],
    station_count: usize,
    trial: usize,
    rng: &mut XorShift,
) -> Vec<Point> {
    if station_count == 0 {
        return Vec::new();
    }
    if planets.is_empty() {
        return vec![Point { x: 500, y: 500 }; station_count];
    }

    let mut centers = initial_centers(planets, station_count, trial, rng);
    let mut assign = vec![0usize; planets.len()];

    for _ in 0..28 {
        for (i, &p) in planets.iter().enumerate() {
            let mut best = 0;
            let mut best_d = dist_to_center(p, centers[0]);
            for (j, &c) in centers.iter().enumerate().skip(1) {
                let d = dist_to_center(p, c);
                if d < best_d {
                    best_d = d;
                    best = j;
                }
            }
            assign[i] = best;
        }

        let mut sx = vec![0.0; station_count];
        let mut sy = vec![0.0; station_count];
        let mut cnt = vec![0usize; station_count];
        for (i, &a) in assign.iter().enumerate() {
            sx[a] += planets[i].x as f64;
            sy[a] += planets[i].y as f64;
            cnt[a] += 1;
        }

        let mut changed = false;
        for j in 0..station_count {
            let next = if cnt[j] > 0 {
                (sx[j] / cnt[j] as f64, sy[j] / cnt[j] as f64)
            } else {
                let p = farthest_planet_from_centers(planets, &centers);
                (p.x as f64, p.y as f64)
            };
            if (centers[j].0 - next.0).abs() > 1.0e-7 || (centers[j].1 - next.1).abs() > 1.0e-7 {
                changed = true;
            }
            centers[j] = next;
        }
        if !changed {
            break;
        }
    }

    centers
        .into_iter()
        .map(|(x, y)| Point {
            x: clamp_coord(x),
            y: clamp_coord(y),
        })
        .collect()
}

fn initial_centers(
    planets: &[Point],
    station_count: usize,
    trial: usize,
    rng: &mut XorShift,
) -> Vec<(f64, f64)> {
    let n = planets.len();
    let first = match trial % 5 {
        0 => 0,
        1 => farthest_from_point(planets, planets[0]),
        2 => closest_to_centroid(planets),
        _ => rng.gen_usize(n),
    };

    let mut centers = vec![(planets[first].x as f64, planets[first].y as f64)];
    let mut selected = vec![false; n];
    selected[first] = true;

    while centers.len() < station_count {
        let idx = if trial % 5 == 3 {
            kmeans_plus_plus_choice(planets, &centers, rng)
        } else {
            let mut best_i = 0;
            let mut best_d = -1.0;
            for (i, &p) in planets.iter().enumerate() {
                if selected[i] && centers.len() < n {
                    continue;
                }
                let d = min_dist_to_centers(p, &centers);
                if d > best_d {
                    best_d = d;
                    best_i = i;
                }
            }
            best_i
        };
        centers.push((planets[idx].x as f64, planets[idx].y as f64));
        selected[idx] = true;
    }
    centers
}

fn farthest_from_point(planets: &[Point], base: Point) -> usize {
    let mut best = 0;
    let mut best_d = -1.0;
    for (i, &p) in planets.iter().enumerate() {
        let d = sq_dist_i(p, base);
        if d > best_d {
            best_d = d;
            best = i;
        }
    }
    best
}

fn closest_to_centroid(planets: &[Point]) -> usize {
    let sx: f64 = planets.iter().map(|p| p.x as f64).sum();
    let sy: f64 = planets.iter().map(|p| p.y as f64).sum();
    let c = (sx / planets.len() as f64, sy / planets.len() as f64);
    let mut best = 0;
    let mut best_d = INF;
    for (i, &p) in planets.iter().enumerate() {
        let d = dist_to_center(p, c);
        if d < best_d {
            best_d = d;
            best = i;
        }
    }
    best
}

fn kmeans_plus_plus_choice(planets: &[Point], centers: &[(f64, f64)], rng: &mut XorShift) -> usize {
    let weights: Vec<f64> = planets
        .iter()
        .map(|&p| min_dist_to_centers(p, centers).max(1.0))
        .collect();
    let total: f64 = weights.iter().sum();
    let mut r = (rng.next_u64() as f64 / u64::MAX as f64) * total;
    for (i, w) in weights.into_iter().enumerate() {
        r -= w;
        if r <= 0.0 {
            return i;
        }
    }
    planets.len() - 1
}

fn farthest_planet_from_centers(planets: &[Point], centers: &[(f64, f64)]) -> Point {
    let mut best = planets[0];
    let mut best_d = -1.0;
    for &p in planets {
        let d = min_dist_to_centers(p, centers);
        if d > best_d {
            best_d = d;
            best = p;
        }
    }
    best
}

fn min_dist_to_centers(p: Point, centers: &[(f64, f64)]) -> f64 {
    centers
        .iter()
        .map(|&c| dist_to_center(p, c))
        .fold(INF, f64::min)
}

fn dist_to_center(p: Point, c: (f64, f64)) -> f64 {
    let dx = p.x as f64 - c.0;
    let dy = p.y as f64 - c.1;
    dx * dx + dy * dy
}

fn clamp_coord(v: f64) -> i32 {
    v.round().clamp(0.0, 1000.0) as i32
}

fn evaluate(planets: &[Point], stations: Vec<Point>) -> Evaluation {
    let (dist, next) = build_shortest_paths(planets, &stations);
    let planet_dist: Vec<Vec<f64>> = (0..planets.len())
        .map(|i| (0..planets.len()).map(|j| dist[i][j]).collect())
        .collect();

    let mut candidates = Vec::new();
    candidates.push(cheapest_insertion_route(&planet_dist));
    candidates.push(nearest_neighbor_route(&planet_dist));
    candidates.push(axis_route(planets, true));
    candidates.push(axis_route(planets, false));
    candidates.push(angle_route(planets));

    let mut best_route = candidates.remove(0);
    improve_route(&mut best_route, &planet_dist);
    let mut best_cost = route_cost(&best_route, &planet_dist);

    for mut route in candidates {
        improve_route(&mut route, &planet_dist);
        let cost = route_cost(&route, &planet_dist);
        if cost < best_cost {
            best_cost = cost;
            best_route = route;
        }
    }

    Evaluation {
        cost: best_cost,
        stations,
        route: best_route,
        next,
    }
}

fn build_shortest_paths(planets: &[Point], stations: &[Point]) -> (Vec<Vec<f64>>, Vec<Vec<usize>>) {
    let n = planets.len();
    let m = stations.len();
    let total = n + m;
    let mut points = Vec::with_capacity(total);
    points.extend_from_slice(planets);
    points.extend_from_slice(stations);

    let mut dist = vec![vec![INF; total]; total];
    let mut next = vec![vec![0usize; total]; total];
    for i in 0..total {
        for j in 0..total {
            if i == j {
                dist[i][j] = 0.0;
                next[i][j] = j;
            } else {
                dist[i][j] = edge_cost(points[i], i < n, points[j], j < n);
                next[i][j] = j;
            }
        }
    }

    for k in 0..total {
        for i in 0..total {
            let dik = dist[i][k];
            if dik >= INF * 0.5 {
                continue;
            }
            for j in 0..total {
                let nd = dik + dist[k][j];
                if nd + 1.0e-9 < dist[i][j] {
                    dist[i][j] = nd;
                    next[i][j] = next[i][k];
                }
            }
        }
    }
    (dist, next)
}

fn edge_cost(a: Point, a_planet: bool, b: Point, b_planet: bool) -> f64 {
    let w = match (a_planet, b_planet) {
        (true, true) => PLANET_WEIGHT * PLANET_WEIGHT,
        (true, false) | (false, true) => PLANET_WEIGHT,
        (false, false) => 1.0,
    };
    w * sq_dist_i(a, b)
}

fn sq_dist_i(a: Point, b: Point) -> f64 {
    let dx = (a.x - b.x) as f64;
    let dy = (a.y - b.y) as f64;
    dx * dx + dy * dy
}

fn cheapest_insertion_route(dist: &[Vec<f64>]) -> Vec<usize> {
    let n = dist.len();
    if n == 0 {
        return Vec::new();
    }
    let mut route = vec![0, 0];
    let mut used = vec![false; n];
    used[0] = true;

    for _ in 1..n {
        let mut best_p = 0;
        let mut best_pos = 0;
        let mut best_inc = INF;
        for p in 1..n {
            if used[p] {
                continue;
            }
            for pos in 0..route.len() - 1 {
                let a = route[pos];
                let b = route[pos + 1];
                let inc = dist[a][p] + dist[p][b] - dist[a][b];
                if inc < best_inc {
                    best_inc = inc;
                    best_p = p;
                    best_pos = pos + 1;
                }
            }
        }
        route.insert(best_pos, best_p);
        used[best_p] = true;
    }
    route
}

fn nearest_neighbor_route(dist: &[Vec<f64>]) -> Vec<usize> {
    let n = dist.len();
    if n == 0 {
        return Vec::new();
    }
    let mut route = Vec::with_capacity(n + 1);
    let mut used = vec![false; n];
    let mut cur = 0;
    route.push(cur);
    used[cur] = true;

    for _ in 1..n {
        let mut best = 0;
        let mut best_d = INF;
        for p in 1..n {
            if !used[p] && dist[cur][p] < best_d {
                best_d = dist[cur][p];
                best = p;
            }
        }
        cur = best;
        used[cur] = true;
        route.push(cur);
    }
    route.push(0);
    route
}

fn axis_route(planets: &[Point], sort_x: bool) -> Vec<usize> {
    if planets.is_empty() {
        return Vec::new();
    }
    let mut ids: Vec<usize> = (1..planets.len()).collect();
    if sort_x {
        ids.sort_by_key(|&i| (planets[i].x, planets[i].y));
    } else {
        ids.sort_by_key(|&i| (planets[i].y, planets[i].x));
    }
    let mut route = Vec::with_capacity(planets.len() + 1);
    route.push(0);
    route.extend(ids);
    route.push(0);
    route
}

fn angle_route(planets: &[Point]) -> Vec<usize> {
    if planets.is_empty() {
        return Vec::new();
    }
    let cx = planets.iter().map(|p| p.x as f64).sum::<f64>() / planets.len() as f64;
    let cy = planets.iter().map(|p| p.y as f64).sum::<f64>() / planets.len() as f64;
    let mut ids: Vec<usize> = (1..planets.len()).collect();
    ids.sort_by(|&i, &j| {
        let ai = (planets[i].y as f64 - cy).atan2(planets[i].x as f64 - cx);
        let aj = (planets[j].y as f64 - cy).atan2(planets[j].x as f64 - cx);
        ai.partial_cmp(&aj).unwrap()
    });
    let mut route = Vec::with_capacity(planets.len() + 1);
    route.push(0);
    route.extend(ids);
    route.push(0);
    route
}

fn improve_route(route: &mut [usize], dist: &[Vec<f64>]) {
    two_opt(route, dist);
    relocate_single(route, dist);
    two_opt(route, dist);
}

fn two_opt(route: &mut [usize], dist: &[Vec<f64>]) {
    if route.len() <= 4 {
        return;
    }
    let last_inner = route.len() - 2;
    let mut improved = true;
    while improved {
        improved = false;
        for i in 1..=last_inner {
            for j in i + 1..=last_inner {
                let a = route[i - 1];
                let b = route[i];
                let c = route[j];
                let d = route[j + 1];
                let delta = dist[a][c] + dist[b][d] - dist[a][b] - dist[c][d];
                if delta < -1.0e-9 {
                    route[i..=j].reverse();
                    improved = true;
                }
            }
        }
    }
}

fn relocate_single(route: &mut [usize], dist: &[Vec<f64>]) {
    if route.len() <= 4 {
        return;
    }
    let len = route.len();
    let last_inner = len - 2;
    let mut improved = true;
    while improved {
        improved = false;
        'outer: for i in 1..=last_inner {
            let p = route[i];
            let prev = route[i - 1];
            let next = route[i + 1];
            let remove_delta = dist[prev][next] - dist[prev][p] - dist[p][next];
            for j in 0..=last_inner {
                if j == i || j + 1 == i {
                    continue;
                }
                let a = route[j];
                let b = route[j + 1];
                let insert_delta = dist[a][p] + dist[p][b] - dist[a][b];
                if remove_delta + insert_delta < -1.0e-9 {
                    if j < i {
                        route[j + 1..=i].rotate_right(1);
                    } else {
                        route[i..=j].rotate_left(1);
                    }
                    improved = true;
                    break 'outer;
                }
            }
        }
    }
}

fn route_cost(route: &[usize], dist: &[Vec<f64>]) -> f64 {
    route.windows(2).map(|w| dist[w[0]][w[1]]).sum()
}

fn improve_stations_for_route(
    planets: &[Point],
    stations: &mut [Point],
    route: &[usize],
    started_at: Instant,
) {
    if stations.is_empty() || planets.is_empty() {
        return;
    }

    let dirs = [
        (1, 0),
        (-1, 0),
        (0, 1),
        (0, -1),
        (1, 1),
        (1, -1),
        (-1, 1),
        (-1, -1),
    ];
    let steps = [64, 32, 16, 8, 4, 2, 1];
    let mut best_cost = fixed_route_cost(planets, stations, route);

    for &step in &steps {
        if started_at.elapsed() >= SOFT_LIMIT {
            break;
        }
        let mut improved = true;
        while improved && started_at.elapsed() < SOFT_LIMIT {
            improved = false;
            for si in 0..stations.len() {
                if started_at.elapsed() >= SOFT_LIMIT {
                    return;
                }
                let original = stations[si];
                let mut local_best = best_cost;
                let mut local_point = original;
                for &(dx, dy) in &dirs {
                    let nx = original.x + dx * step;
                    let ny = original.y + dy * step;
                    if !(0..=1000).contains(&nx) || !(0..=1000).contains(&ny) {
                        continue;
                    }
                    stations[si] = Point { x: nx, y: ny };
                    let cost = fixed_route_cost(planets, stations, route);
                    if cost + 1.0e-9 < local_best {
                        local_best = cost;
                        local_point = stations[si];
                    }
                }
                stations[si] = local_point;
                if local_best + 1.0e-9 < best_cost {
                    best_cost = local_best;
                    improved = true;
                }
            }
        }
    }
}

fn optimize_stations_for_path(planets: &[Point], stations: &[Point], path: &[usize]) -> Vec<Point> {
    let n = planets.len();
    let m = stations.len();
    if m == 0 {
        return Vec::new();
    }

    let mut a = vec![vec![0.0; m]; m];
    let mut bx = vec![0.0; m];
    let mut by = vec![0.0; m];

    for w in path.windows(2) {
        let u = w[0];
        let v = w[1];
        if u == v || (u < n && v < n) {
            continue;
        }
        let coeff = if u >= n && v >= n { 1.0 } else { PLANET_WEIGHT };

        match (u < n, v < n) {
            (false, false) => {
                let su = u - n;
                let sv = v - n;
                a[su][su] += coeff;
                a[sv][sv] += coeff;
                a[su][sv] -= coeff;
                a[sv][su] -= coeff;
            }
            (false, true) => {
                let su = u - n;
                a[su][su] += coeff;
                bx[su] += coeff * planets[v].x as f64;
                by[su] += coeff * planets[v].y as f64;
            }
            (true, false) => {
                let sv = v - n;
                a[sv][sv] += coeff;
                bx[sv] += coeff * planets[u].x as f64;
                by[sv] += coeff * planets[u].y as f64;
            }
            (true, true) => unreachable!(),
        }
    }

    for i in 0..m {
        if a[i][i].abs() < 1.0e-9 {
            a[i][i] = 1.0;
            bx[i] = stations[i].x as f64;
            by[i] = stations[i].y as f64;
        }
    }

    let xs = solve_linear(a.clone(), bx);
    let ys = solve_linear(a, by);
    (0..m)
        .map(|i| Point {
            x: clamp_coord(xs[i]),
            y: clamp_coord(ys[i]),
        })
        .collect()
}

fn solve_linear(mut a: Vec<Vec<f64>>, mut b: Vec<f64>) -> Vec<f64> {
    let n = b.len();
    for col in 0..n {
        let mut pivot = col;
        for row in col + 1..n {
            if a[row][col].abs() > a[pivot][col].abs() {
                pivot = row;
            }
        }
        if a[pivot][col].abs() < 1.0e-9 {
            continue;
        }
        a.swap(col, pivot);
        b.swap(col, pivot);

        let div = a[col][col];
        for j in col..n {
            a[col][j] /= div;
        }
        b[col] /= div;

        for row in 0..n {
            if row == col {
                continue;
            }
            let factor = a[row][col];
            if factor.abs() < 1.0e-12 {
                continue;
            }
            for j in col..n {
                a[row][j] -= factor * a[col][j];
            }
            b[row] -= factor * b[col];
        }
    }
    b
}

fn fixed_route_cost(planets: &[Point], stations: &[Point], route: &[usize]) -> f64 {
    let (dist, _) = build_shortest_paths(planets, stations);
    route_cost(route, &dist)
}

fn expand_route(planet_count: usize, planet_route: &[usize], next: &[Vec<usize>]) -> Vec<usize> {
    if planet_route.is_empty() {
        return Vec::new();
    }
    let mut path = vec![planet_route[0]];
    let total = next.len();
    for w in planet_route.windows(2) {
        let target = w[1];
        let mut cur = w[0];
        let mut guard = 0;
        while cur != target && guard <= total + planet_count + 5 {
            cur = next[cur][target];
            path.push(cur);
            guard += 1;
        }
    }
    path
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn station_shortcut_is_used() {
        let planets = vec![Point { x: 0, y: 0 }, Point { x: 200, y: 0 }];
        let stations = vec![Point { x: 100, y: 0 }];
        let (dist, next) = build_shortest_paths(&planets, &stations);
        assert!((dist[0][1] - 100_000.0).abs() < 1.0e-6);
        let path = expand_route(planets.len(), &[0, 1, 0], &next);
        assert_eq!(path[0], 0);
        assert_eq!(*path.last().unwrap(), 0);
        assert!(path.contains(&2));
    }

    #[test]
    fn generated_answer_visits_all_planets() {
        let planets = vec![
            Point { x: 100, y: 100 },
            Point { x: 0, y: 0 },
            Point { x: 0, y: 100 },
        ];
        let answer = solve(&planets, 4);
        assert_eq!(answer.path.first(), Some(&0));
        assert_eq!(answer.path.last(), Some(&0));
        for i in 0..planets.len() {
            assert!(answer.path.contains(&i));
        }
        assert_eq!(answer.stations.len(), 4);
        assert!(answer.path.len() <= 100_000);
    }
}