#1169523 (Rust) No.5007 Steiner Space Travel

• 問題: No.5007 Steiner Space Travel
• ユーザー: assy1028
• 提出日時: 2026-06-21 13:21:31
• 言語: Rust (1.94.0 + proconio + num + itertools)
• 結果: AC
• 実行時間: 943 ms / 1,000 ms
• コード長: 21,422 bytes
• コンパイル時間: 2,642 ms
• コンパイル使用メモリ: 212,676 KB
• 実行使用メモリ: 6,400 KB
• スコア: 8,900,924
• 最終ジャッジ日時: 2026-06-21 13:22:08
• 合計ジャッジ時間: 33,498 ms
• ジャッジサーバーID （参考情報）: judge3_1 / judge1_0

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

ソースコード

use std::cmp::min;
use std::io::{self, Read, Write};
use std::time::Instant;

const ALPHA: i64 = 5;
const TIME_LIMIT: f64 = 0.94;
const INF: i64 = 1_i64 << 60;

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

impl Point {
    fn dist2(self, other: Point) -> i64 {
        let dx = self.x - other.x;
        let dy = self.y - other.y;
        dx * dx + dy * dy
    }
}

struct Scanner {
    input: Vec<u8>,
    index: usize,
}

impl Scanner {
    fn new() -> Self {
        let mut input = String::new();
        io::stdin().read_to_string(&mut input).unwrap();
        Self {
            input: input.into_bytes(),
            index: 0,
        }
    }

    fn next_i64(&mut self) -> i64 {
        while self.index < self.input.len() && self.input[self.index].is_ascii_whitespace() {
            self.index += 1;
        }
        let mut sign = 1;
        if self.index < self.input.len() && self.input[self.index] == b'-' {
            sign = -1;
            self.index += 1;
        }
        let mut value = 0_i64;
        while self.index < self.input.len() && !self.input[self.index].is_ascii_whitespace() {
            value = value * 10 + (self.input[self.index] - b'0') as i64;
            self.index += 1;
        }
        value * sign
    }
}

#[derive(Clone)]
struct XorShift {
    state: u64,
}

impl XorShift {
    fn new(seed: u64) -> Self {
        Self { state: seed.max(1) }
    }

    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 {
        (self.next_u64() % upper as u64) as usize
    }

    fn gen_f64(&mut self) -> f64 {
        const DEN: f64 = (1_u64 << 53) as f64;
        ((self.next_u64() >> 11) as f64) / DEN
    }
}

#[derive(Clone)]
struct Metric {
    n: usize,
    m: usize,
    planet_dist: Vec<i64>,
    choice: Vec<(i32, i32)>,
    station_next: Vec<usize>,
}

impl Metric {
    fn choice(&self, a: usize, b: usize) -> (i32, i32) {
        self.choice[a * self.n + b]
    }

    fn snext(&self, a: usize, b: usize) -> usize {
        self.station_next[a * self.m + b]
    }
}

fn elapsed(start: Instant) -> f64 {
    start.elapsed().as_secs_f64()
}

fn clamp_coord(v: i64) -> i64 {
    v.clamp(0, 1000)
}

fn compute_metric(planets: &[Point], stations: &[Point]) -> Metric {
    let n = planets.len();
    let m = stations.len();
    let mut station_dist = vec![0_i64; m * m];
    let mut station_next = vec![0_usize; m * m];

    for i in 0..m {
        for j in 0..m {
            station_dist[i * m + j] = stations[i].dist2(stations[j]);
            station_next[i * m + j] = j;
        }
    }
    for k in 0..m {
        for i in 0..m {
            let dik = station_dist[i * m + k];
            for j in 0..m {
                let v = dik + station_dist[k * m + j];
                if v < station_dist[i * m + j] {
                    station_dist[i * m + j] = v;
                    station_next[i * m + j] = station_next[i * m + k];
                }
            }
        }
    }

    let mut planet_station = vec![0_i64; n * m];
    for i in 0..n {
        for s in 0..m {
            planet_station[i * m + s] = ALPHA * planets[i].dist2(stations[s]);
        }
    }

    let mut planet_dist = vec![0_i64; n * n];
    let mut choice = vec![(-1_i32, -1_i32); n * n];
    for i in 0..n {
        for j in 0..n {
            if i == j {
                continue;
            }
            let mut best = ALPHA * ALPHA * planets[i].dist2(planets[j]);
            let mut best_choice = (-1_i32, -1_i32);
            for a in 0..m {
                let start_cost = planet_station[i * m + a];
                for b in 0..m {
                    let v = start_cost + station_dist[a * m + b] + planet_station[j * m + b];
                    if v < best {
                        best = v;
                        best_choice = (a as i32, b as i32);
                    }
                }
            }
            planet_dist[i * n + j] = best;
            choice[i * n + j] = best_choice;
        }
    }

    Metric {
        n,
        m,
        planet_dist,
        choice,
        station_next,
    }
}

fn compute_direct_metric(planets: &[Point]) -> Vec<i64> {
    let n = planets.len();
    let mut dist = vec![0_i64; n * n];
    for i in 0..n {
        for j in 0..n {
            if i != j {
                dist[i * n + j] = ALPHA * ALPHA * planets[i].dist2(planets[j]);
            }
        }
    }
    dist
}

fn route_cost(route: &[usize], dist: &[i64], n: usize) -> i64 {
    let mut cost = 0_i64;
    for i in 0..route.len() {
        cost += dist[route[i] * n + route[(i + 1) % route.len()]];
    }
    cost
}

fn nearest_route(dist: &[i64], n: usize, first_rank: usize) -> Vec<usize> {
    let mut used = vec![false; n];
    let mut route = Vec::with_capacity(n);
    route.push(0);
    used[0] = true;
    let mut cur = 0;
    for step in 0..n - 1 {
        let mut cand = Vec::with_capacity(n - step - 1);
        for j in 0..n {
            if !used[j] {
                cand.push((dist[cur * n + j], j));
            }
        }
        cand.sort_unstable();
        let rank = if step == 0 {
            min(first_rank, cand.len() - 1)
        } else {
            0
        };
        let next = cand[rank].1;
        used[next] = true;
        route.push(next);
        cur = next;
    }
    route
}

fn cheapest_insertion_route(dist: &[i64], n: usize) -> Vec<usize> {
    let mut route = vec![0];
    let mut unused = vec![true; n];
    unused[0] = false;

    let mut farthest = 1;
    for i in 2..n {
        if dist[i] > dist[farthest] {
            farthest = i;
        }
    }
    route.push(farthest);
    unused[farthest] = false;

    while route.len() < n {
        let mut best = (INF, 0_usize, 0_usize);
        for x in 1..n {
            if !unused[x] {
                continue;
            }
            for pos in 0..route.len() {
                let a = route[pos];
                let b = route[(pos + 1) % route.len()];
                let inc = dist[a * n + x] + dist[x * n + b] - dist[a * n + b];
                if inc < best.0 {
                    best = (inc, x, pos + 1);
                }
            }
        }
        route.insert(best.2, best.1);
        unused[best.1] = false;
    }
    let z = route.iter().position(|&v| v == 0).unwrap();
    route.rotate_left(z);
    route
}

fn random_route(n: usize, rng: &mut XorShift) -> Vec<usize> {
    let mut route = (0..n).collect::<Vec<_>>();
    for i in 1..n {
        let j = 1 + rng.gen_usize(n - 1);
        route.swap(i, j);
    }
    route
}

fn improve_route(route: &mut Vec<usize>, dist: &[i64], n: usize, start: Instant, limit: f64) {
    let mut iter = 0;
    loop {
        iter += 1;
        if iter & 7 == 0 && elapsed(start) > limit {
            break;
        }
        let mut best_delta = 0_i64;
        let mut op_kind = 0_i32;
        let mut op_i = 0_usize;
        let mut op_j = 0_usize;

        for i in 1..n - 1 {
            let a = route[i - 1];
            let b = route[i];
            for k in i + 1..n {
                let c = route[k];
                let d = route[(k + 1) % n];
                let delta = dist[a * n + c] + dist[b * n + d] - dist[a * n + b] - dist[c * n + d];
                if delta < best_delta {
                    best_delta = delta;
                    op_kind = 0;
                    op_i = i;
                    op_j = k;
                }
            }
        }

        for len in 1..=3 {
            for i in 1..=n - len {
                let prev = route[i - 1];
                let first = route[i];
                let last = route[i + len - 1];
                let next = route[(i + len) % n];
                let remove_delta =
                    dist[prev * n + next] - dist[prev * n + first] - dist[last * n + next];
                for j in 1..=n {
                    if i <= j && j <= i + len {
                        continue;
                    }
                    let jj = if j == n { 0 } else { j };
                    let a = route[j - 1];
                    let b = route[jj];
                    let delta =
                        remove_delta + dist[a * n + first] + dist[last * n + b] - dist[a * n + b];
                    if delta < best_delta {
                        best_delta = delta;
                        op_kind = len as i32;
                        op_i = i;
                        op_j = j;
                    }
                }
            }
        }

        if best_delta >= 0 {
            break;
        }

        if op_kind == 0 {
            route[op_i..=op_j].reverse();
        } else {
            let len = op_kind as usize;
            let seg = route[op_i..op_i + len].to_vec();
            route.drain(op_i..op_i + len);
            let mut pos = op_j;
            if pos > op_i {
                pos -= len;
            }
            route.splice(pos..pos, seg);
        }
    }
}

fn solve_tsp(
    dist: &[i64],
    n: usize,
    rng: &mut XorShift,
    start: Instant,
    limit: f64,
    tries: usize,
) -> (i64, Vec<usize>) {
    let mut best_cost = INF;
    let mut best_route = Vec::new();
    for t in 0..tries {
        if elapsed(start) > limit {
            break;
        }
        let mut route = if t < 6 {
            nearest_route(dist, n, t)
        } else if t == 6 {
            cheapest_insertion_route(dist, n)
        } else {
            random_route(n, rng)
        };
        improve_route(&mut route, dist, n, start, limit);
        let cost = route_cost(&route, dist, n);
        if cost < best_cost {
            best_cost = cost;
            best_route = route;
        }
    }
    if best_route.is_empty() {
        let mut route = nearest_route(dist, n, 0);
        improve_route(&mut route, dist, n, start, limit);
        best_cost = route_cost(&route, dist, n);
        best_route = route;
    }
    (best_cost, best_route)
}

fn kmeans_init(planets: &[Point], m: usize, seed: u64) -> Vec<Point> {
    let n = planets.len();
    let mut rng = XorShift::new(seed);
    let mut centers = Vec::<(f64, f64)>::with_capacity(m);
    let first = rng.gen_usize(n);
    centers.push((planets[first].x as f64, planets[first].y as f64));

    while centers.len() < m {
        let mut dists = vec![0_f64; n];
        let mut sum = 0_f64;
        for (i, p) in planets.iter().enumerate() {
            let mut best = f64::INFINITY;
            for &(cx, cy) in &centers {
                let dx = p.x as f64 - cx;
                let dy = p.y as f64 - cy;
                best = best.min(dx * dx + dy * dy);
            }
            dists[i] = best;
            sum += best;
        }
        let idx = if sum <= 1e-9 {
            rng.gen_usize(n)
        } else {
            let mut target = rng.gen_f64() * sum;
            let mut selected = n - 1;
            for (i, &d) in dists.iter().enumerate() {
                target -= d;
                if target <= 0.0 {
                    selected = i;
                    break;
                }
            }
            selected
        };
        centers.push((planets[idx].x as f64, planets[idx].y as f64));
    }

    for _ in 0..20 {
        let mut sx = vec![0_f64; m];
        let mut sy = vec![0_f64; m];
        let mut cnt = vec![0_usize; m];
        for p in planets {
            let mut best = 0;
            let mut best_d = f64::INFINITY;
            for (i, &(cx, cy)) in centers.iter().enumerate() {
                let dx = p.x as f64 - cx;
                let dy = p.y as f64 - cy;
                let d = dx * dx + dy * dy;
                if d < best_d {
                    best_d = d;
                    best = i;
                }
            }
            sx[best] += p.x as f64;
            sy[best] += p.y as f64;
            cnt[best] += 1;
        }
        for i in 0..m {
            if cnt[i] > 0 {
                centers[i] = (sx[i] / cnt[i] as f64, sy[i] / cnt[i] as f64);
            } else {
                let p = planets[rng.gen_usize(n)];
                centers[i] = (p.x as f64, p.y as f64);
            }
        }
    }

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

fn farthest_init(planets: &[Point], m: usize) -> Vec<Point> {
    let mut centers = vec![planets[0]];
    while centers.len() < m {
        let mut best_i = 0;
        let mut best_d = -1_i64;
        for (i, &p) in planets.iter().enumerate() {
            let d = centers.iter().map(|&c| p.dist2(c)).min().unwrap();
            if d > best_d {
                best_d = d;
                best_i = i;
            }
        }
        centers.push(planets[best_i]);
    }
    centers
}

fn octagon_init(planets: &[Point], m: usize) -> Vec<Point> {
    let n = planets.len() as f64;
    let cx = planets.iter().map(|p| p.x as f64).sum::<f64>() / n;
    let cy = planets.iter().map(|p| p.y as f64).sum::<f64>() / n;
    let mut radius = 0_f64;
    for p in planets {
        let dx = p.x as f64 - cx;
        let dy = p.y as f64 - cy;
        radius += (dx * dx + dy * dy).sqrt();
    }
    radius = (radius / n).clamp(120.0, 320.0);
    (0..m)
        .map(|i| {
            let th = std::f64::consts::TAU * i as f64 / m as f64;
            Point {
                x: clamp_coord((cx + radius * th.cos()).round() as i64),
                y: clamp_coord((cy + radius * th.sin()).round() as i64),
            }
        })
        .collect()
}

fn midpoint_init(planets: &[Point], m: usize, start: Instant) -> Vec<Point> {
    let n = planets.len();
    let direct = compute_direct_metric(planets);
    let mut rng = XorShift::new(1_234_567);
    let (_, route) = solve_tsp(&direct, n, &mut rng, start, 0.20, 4);
    let mut edges = Vec::with_capacity(n);
    for i in 0..n {
        let a = route[i];
        let b = route[(i + 1) % n];
        edges.push((planets[a].dist2(planets[b]), a, b));
    }
    edges.sort_unstable_by(|a, b| b.0.cmp(&a.0));
    edges
        .into_iter()
        .take(m)
        .map(|(_, a, b)| Point {
            x: (planets[a].x + planets[b].x + 1) / 2,
            y: (planets[a].y + planets[b].y + 1) / 2,
        })
        .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 pivot != col {
            a.swap(pivot, col);
            b.swap(pivot, col);
        }
        let div = a[col][col];
        if div.abs() < 1e-12 {
            continue;
        }
        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() < 1e-15 {
                continue;
            }
            for j in col..n {
                a[row][j] -= factor * a[col][j];
            }
            b[row] -= factor * b[col];
        }
    }
    b
}

fn optimize_stations(
    planets: &[Point],
    stations: &[Point],
    route: &[usize],
    metric: &Metric,
) -> Vec<Point> {
    let m = stations.len();
    let mut mat = vec![vec![0_f64; m]; m];
    let mut bx = vec![0_f64; m];
    let mut by = vec![0_f64; m];

    for i in 0..route.len() {
        let p = route[i];
        let q = route[(i + 1) % route.len()];
        let (a_raw, b_raw) = metric.choice(p, q);
        if a_raw < 0 {
            continue;
        }
        let a = a_raw as usize;
        let b = b_raw as usize;
        add_station_planet(a, planets[p], 5.0, &mut mat, &mut bx, &mut by);
        let mut cur = a;
        while cur != b {
            let next = metric.snext(cur, b);
            add_station_station(cur, next, 1.0, &mut mat);
            cur = next;
        }
        add_station_planet(b, planets[q], 5.0, &mut mat, &mut bx, &mut by);
    }

    for i in 0..m {
        mat[i][i] += 1e-6;
        bx[i] += 1e-6 * stations[i].x as f64;
        by[i] += 1e-6 * stations[i].y as f64;
    }

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

fn add_station_planet(
    s: usize,
    p: Point,
    w: f64,
    mat: &mut [Vec<f64>],
    bx: &mut [f64],
    by: &mut [f64],
) {
    mat[s][s] += w;
    bx[s] += w * p.x as f64;
    by[s] += w * p.y as f64;
}

fn add_station_station(a: usize, b: usize, w: f64, mat: &mut [Vec<f64>]) {
    if a == b {
        return;
    }
    mat[a][a] += w;
    mat[b][b] += w;
    mat[a][b] -= w;
    mat[b][a] -= w;
}

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

fn try_candidate(
    planets: &[Point],
    mut stations: Vec<Point>,
    best: &mut Option<Answer>,
    rng: &mut XorShift,
    start: Instant,
    limit: f64,
    tsp_tries: usize,
    refine_count: usize,
) {
    for _ in 0..refine_count {
        if elapsed(start) > limit {
            break;
        }
        let metric = compute_metric(planets, &stations);
        let (cost, route) = solve_tsp(
            &metric.planet_dist,
            planets.len(),
            rng,
            start,
            limit,
            tsp_tries,
        );
        if best.as_ref().map_or(true, |ans| cost < ans.cost) {
            *best = Some(Answer {
                cost,
                stations: stations.clone(),
                route: route.clone(),
            });
        }
        let next_stations = optimize_stations(planets, &stations, &route, &metric);
        if next_stations == stations {
            break;
        }
        stations = next_stations;
    }
}

fn push_node(nodes: &mut Vec<(i32, usize)>, node: (i32, usize)) {
    if nodes.last().copied() != Some(node) {
        nodes.push(node);
    }
}

fn expand_route(answer: &Answer, planets: &[Point]) -> Vec<(i32, usize)> {
    let metric = compute_metric(planets, &answer.stations);
    let n = planets.len();
    let mut nodes = Vec::new();
    push_node(&mut nodes, (1, answer.route[0]));
    for i in 0..n {
        let p = answer.route[i];
        let q = answer.route[(i + 1) % n];
        let (a_raw, b_raw) = metric.choice(p, q);
        if a_raw >= 0 {
            let a = a_raw as usize;
            let b = b_raw as usize;
            push_node(&mut nodes, (2, a));
            let mut cur = a;
            while cur != b {
                cur = metric.snext(cur, b);
                push_node(&mut nodes, (2, cur));
            }
        }
        push_node(&mut nodes, (1, q));
    }
    nodes
}

fn main() {
    let start = Instant::now();
    let mut sc = Scanner::new();
    let n = sc.next_i64() as usize;
    let m = sc.next_i64() as usize;
    let mut planets = Vec::with_capacity(n);
    for _ in 0..n {
        let x = sc.next_i64();
        let y = sc.next_i64();
        planets.push(Point { x, y });
    }

    let mut rng = XorShift::new(0x9E37_79B9_7F4A_7C15);
    let mut best: Option<Answer> = None;

    let mut presets = Vec::new();
    for seed in [0_u64, 4, 13, 29] {
        presets.push(kmeans_init(&planets, m, seed));
    }
    presets.push(farthest_init(&planets, m));
    presets.push(octagon_init(&planets, m));
    presets.push(midpoint_init(&planets, m, start));

    for stations in presets {
        try_candidate(
            &planets, stations, &mut best, &mut rng, start, TIME_LIMIT, 8, 5,
        );
    }

    let mut seed = 0_u64;
    while elapsed(start) < TIME_LIMIT * 0.80 {
        let stations = kmeans_init(&planets, m, 1000 + seed * 7919);
        try_candidate(
            &planets, stations, &mut best, &mut rng, start, TIME_LIMIT, 6, 5,
        );
        seed += 1;
    }

    if let Some(ans) = best.clone() {
        try_candidate(
            &planets,
            ans.stations,
            &mut best,
            &mut rng,
            start,
            TIME_LIMIT,
            20,
            6,
        );
    }

    let answer = best.unwrap_or_else(|| {
        let stations = kmeans_init(&planets, m, 1);
        let metric = compute_metric(&planets, &stations);
        let (cost, route) = solve_tsp(&metric.planet_dist, n, &mut rng, start, TIME_LIMIT, 4);
        Answer {
            cost,
            stations,
            route,
        }
    });

    let nodes = expand_route(&answer, &planets);
    let mut out = String::new();
    for p in &answer.stations {
        out.push_str(&format!("{} {}\n", p.x, p.y));
    }
    out.push_str(&format!("{}\n", nodes.len()));
    for (t, idx) in nodes {
        out.push_str(&format!("{} {}\n", t, idx + 1));
    }
    let mut writer = io::BufWriter::new(io::stdout());
    writer.write_all(out.as_bytes()).unwrap();
}

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

    #[test]
    fn linear_solver_handles_diagonal() {
        let a = vec![vec![2.0, 0.0], vec![0.0, 4.0]];
        let b = vec![6.0, 20.0];
        let x = solve_linear(a, b);
        assert!((x[0] - 3.0).abs() < 1e-9);
        assert!((x[1] - 5.0).abs() < 1e-9);
    }
}