Submission #42195967

---

Source Code Expand

#![allow(non_snake_case)]

use itertools::Itertools;
use proconio::{input, marker::Usize1};
use rustc_hash::FxHashMap;
use std::{cmp::Reverse, collections::BinaryHeap, time::Instant};

use rand::Rng;

const CLIMBING_TEMPERETURE: f64 = 0.;
const RNG_SEED: u128 = 20;

const N: usize = 100;
const INF: usize = 1 << 60;
#[allow(dead_code)]
const TIME_LIMIT: f64 = 1.9;

fn main() {
    let timer = TimeKeeper::new(TIME_LIMIT);
    let input = Input::from_stdin();
    let solution = solve(&input, &timer);
    println!("{}", solution);
    eprintln!("time: {:.3}[s]", timer.instant.elapsed().as_secs_f64());
}

fn choose_powers(
    input: &Input,
    state: &mut State,
    dists_from_house_to_station: &Vec<Vec<usize>>,
    dists: &Vec<Vec<usize>>,
    path: &Vec<Vec<Vec<usize>>>,
) {
    let mut power_heap = BinaryHeap::new();

    state.clear();

    for h in 0..input.K {
        let mut min_dist = INF;
        let mut min_station = 0;
        for s in 0..N {
            if !state.use_station[s] {
                continue;
            }
            if dists_from_house_to_station[h][s] < min_dist {
                min_dist = dists_from_house_to_station[h][s];
                min_station = s;
            }
        }
        power_heap.push((min_dist, min_station, h));
    }

    while let Some((dist, station, house)) = power_heap.pop() {
        if state.coverd[house] {
            continue;
        }
        for house in 0..input.K {
            if dists_from_house_to_station[house][station] <= dist {
                state.coverd[house] = true;
            }
        }
        state.powers[station] = dist;
    }

    let mut used = vec![false; N];
    used[0] = true;
    for i in 0..N {
        if state.powers[i] > 0 {
            used[i] = true;
        }
    }

    let edge_used = spaning_tree(&input, &used, &dists, &path);
    for i in 0..input.M {
        state.switches[i] = edge_used[i];
    }
    state.use_station = used;

    state.compute_score(&input);
}

fn solve(input: &Input, timer: &TimeKeeper) -> State {
    let mut state = State::new(&input);
    let (dists, path) = warshall_floyd(&input.graph);

    let dists_from_house_to_station = dist_house_to_station(&input);
    for i in 0..N {
        state.use_station[i] = true;
    }

    eprintln!("time: {:.3}[s]", timer.instant.elapsed().as_secs_f64());

    choose_powers(
        input,
        &mut state,
        &dists_from_house_to_station,
        &dists,
        &path,
    );

    let temperature_scheduler = ClimbingTemperatureSchedule {};

    simulated_annealing(
        &mut state,
        &timer,
        &temperature_scheduler,
        &input,
        &dists_from_house_to_station,
        &dists,
        &path,
    );

    eprintln!("score: {}", state.score);
    state
}

fn spaning_tree(
    input: &Input,
    used: &Vec<bool>,
    dists: &Vec<Vec<usize>>,
    path: &Vec<Vec<Vec<usize>>>,
) -> Vec<bool> {
    let mut uf = UnionFind::new(N);
    let mut edges_heap = BinaryHeap::new();
    let mut edge_used = vec![false; input.M];
    for i in 0..N {
        if !used[i] {
            continue;
        }
        for j in i + 1..N {
            if !used[j] {
                continue;
            }

            let mut edge_ids = vec![];
            for k in 0..path[i][j].len() - 1 {
                let from = path[i][j][k];
                let to = path[i][j][k + 1];
                let edge_id = input.edge_dicts[&(from, to)];
                edge_ids.push(edge_id);
            }

            edges_heap.push(Reverse((dists[i][j], edge_ids, path[i][j].clone())));
        }
    }

    while let Some(Reverse((_cost, edge_ids, path))) = edges_heap.pop() {
        let from = path[0];
        let to = path[path.len() - 1];
        if uf.is_same(from, to) {
            continue;
        } else {
            uf.unite(from, to);
            for edge_id in edge_ids {
                edge_used[edge_id] = true;
            }
        }
    }
    edge_used
}

fn dist_house_to_station(input: &Input) -> Vec<Vec<usize>> {
    let stations = &input.stations;
    let houses = &input.houses;
    let mut dists = vec![vec![INF; N]; input.K];
    for s in 0..N {
        for h in 0..input.K {
            dists[h][s] = stations[s].dist_ceil(&houses[h]);
            // dists[h][s] = stations[s].dist(&houses[h]);
        }
    }
    dists
}

// ワーシャルフロイド法(経路復元有り)
#[allow(dead_code)]
fn warshall_floyd(graph: &Vec<Vec<Edge>>) -> (Vec<Vec<usize>>, Vec<Vec<Vec<usize>>>) {
    let mut dist = vec![vec![INF; N]; N];
    // どの頂点を経由してiからjに行くかを表す
    let mut prev = vec![vec![INF; N]; N];
    for i in 0..N {
        dist[i][i] = 0;
    }
    for i in 0..N {
        for edge in graph[i].iter() {
            dist[edge.from][edge.to] = edge.cost;
            prev[edge.from][edge.to] = edge.from;
        }
    }
    for k in 0..N {
        for i in 0..N {
            for j in 0..N {
                if dist[i][j] > dist[i][k] + dist[k][j] {
                    dist[i][j] = dist[i][k] + dist[k][j];
                    prev[i][j] = prev[k][j];
                }
            }
        }
    }

    let mut path = vec![vec![vec![]; N]; N];
    for i in 0..N {
        for j in 0..N {
            path[i][j] = restore_path(&prev, i, j);
        }
    }

    (dist, path)
}

// prevをもとに経路復元
#[allow(dead_code)]
fn restore_path(prev: &Vec<Vec<usize>>, s: usize, t: usize) -> Vec<usize> {
    let mut path = vec![];
    let mut cur = t;
    while cur != s {
        path.push(cur);
        cur = prev[s][cur];
    }
    path.push(s);
    path.reverse();
    path
}

#[derive(Debug, Clone)]
struct Point {
    x: isize,
    y: isize,
}

impl Point {
    fn new(x: isize, y: isize) -> Self {
        Self { x, y }
    }

    #[allow(dead_code)]
    fn dist(&self, other: &Self) -> usize {
        let dx = self.x - other.x;
        let dy = self.y - other.y;
        ((dx * dx + dy * dy) as f64).sqrt().round() as usize
    }
    // Pのほうはroundではだめ？
    fn dist_ceil(&self, other: &Self) -> usize {
        let dx = self.x - other.x;
        let dy = self.y - other.y;
        ((dx * dx + dy * dy) as f64).sqrt().ceil() as usize
    }
}

#[derive(Debug, Clone)]
struct Edge {
    from: usize,
    to: usize,
    cost: usize,
}

impl Edge {
    fn new(from: usize, to: usize, cost: usize) -> Self {
        Self { from, to, cost }
    }
}

#[derive(Debug, Clone)]
#[allow(dead_code)]
struct Input {
    M: usize,
    K: usize,
    stations: Vec<Point>,
    UVW: Vec<(usize, usize, usize)>,
    edges: Vec<Edge>,
    houses: Vec<Point>,
    graph: Vec<Vec<Edge>>,
    edge_dicts: FxHashMap<(usize, usize), usize>,
}

impl Input {
    fn from_stdin() -> Self {
        input! {
            _N: usize,
            M: usize,
            K: usize,
            XY_: [(isize, isize); N],
            UVW: [(Usize1, Usize1, usize); M],
            AB_: [(isize, isize); K],
        }

        let stations = XY_.into_iter().map(|(x, y)| Point::new(x, y)).collect();
        let houses = AB_.into_iter().map(|(x, y)| Point::new(x, y)).collect();
        let edges: Vec<Edge> = UVW
            .clone()
            .into_iter()
            .map(|(u, v, w)| Edge::new(u, v, w))
            .collect();

        let mut graph = vec![vec![]; N];
        for &(u, v, cost) in UVW.iter() {
            graph[u].push(Edge::new(u, v, cost));
            graph[v].push(Edge::new(v, u, cost));
        }

        let mut edge_dicts = FxHashMap::default();
        for (i, e) in edges.iter().enumerate() {
            edge_dicts.entry((e.from, e.to)).or_insert(i);
            edge_dicts.entry((e.to, e.from)).or_insert(i);
        }

        Self {
            M,
            K,
            stations,
            UVW,
            edges,
            houses,
            graph,
            edge_dicts,
        }
    }
}

/// 時間計測を行うstruct
#[derive(Debug, Clone)]
struct TimeKeeper {
    instant: Instant,
    time_limit: f64,
}

impl TimeKeeper {
    fn new(time_limit: f64) -> Self {
        Self {
            instant: Instant::now(),
            time_limit,
        }
    }

    fn intime(&self) -> bool {
        self.instant.elapsed().as_secs_f64() < self.time_limit
    }
}

/// 最大化問題としたときの確率(最小化のときはマイナスしとけばよい？)
#[inline]
fn probability(new_score: usize, old_score: usize, temperature: f64) -> f64 {
    if temperature == CLIMBING_TEMPERETURE {
        0.
    } else {
        f64::exp((new_score as f64 - old_score as f64) / temperature)
    }
}

/// 温度管理トレイト
trait TemperatureScheduler {
    fn get_temperature(&self, timekeeper: &TimeKeeper) -> f64;
}

/// 指数スケジューリング
#[derive(Debug, Clone)]
struct ExponentialTemperatureSchedule {
    initial_temperature: f64,
    end_temperature: f64,
}

impl ExponentialTemperatureSchedule {
    fn new(initial_temperature: f64, end_temperature: f64) -> Self {
        Self {
            initial_temperature,
            end_temperature,
        }
    }
}

impl TemperatureScheduler for ExponentialTemperatureSchedule {
    fn get_temperature(&self, timekeeper: &TimeKeeper) -> f64 {
        let t = timekeeper.instant.elapsed().as_secs_f64() / timekeeper.time_limit;
        self.initial_temperature.powf(1.0 - t) * self.end_temperature.powf(t)
    }
}

/// 山登り用温度管理
struct ClimbingTemperatureSchedule {}

impl TemperatureScheduler for ClimbingTemperatureSchedule {
    fn get_temperature(&self, _: &TimeKeeper) -> f64 {
        CLIMBING_TEMPERETURE
    }
}

trait AnnealingState {
    fn get_score(&self) -> usize;
    fn rollback(&mut self, diff: &DiffState);
}

#[derive(Debug, Clone)]
struct State {
    score: usize,
    powers: Vec<usize>,
    switches: Vec<bool>,
    coverd: Vec<bool>,
    use_station: Vec<bool>,
}

impl State {
    fn new(input: &Input) -> Self {
        Self {
            score: 0,
            powers: vec![0; N],
            switches: vec![false; input.M],
            coverd: vec![false; input.K],
            use_station: vec![false; N],
        }
    }

    fn clear(&mut self) {
        self.score = 0;
        self.powers = vec![0; N];
        self.switches = vec![false; self.switches.len()];
        self.coverd = vec![false; self.coverd.len()];
    }

    fn compute_score(&mut self, input: &Input) {
        if self.coverd.iter().all(|x| *x) {
            let mut S = 0;
            for &p in self.powers.iter() {
                if p > 5000 {
                    self.score = 0;
                    return;
                }
                S += p * p;
            }
            for i in 0..input.M {
                if self.switches[i] {
                    S += input.edges[i].cost;
                }
            }
            self.score =
                (1_000_000. * (1. + 100_000_000. / (S as f64 + 10_000_000.))).round() as usize;
        } else {
            let mut n = 0;
            for &c in self.coverd.iter() {
                if c {
                    n += 1;
                }
            }
            self.score = (1_000_000. * ((1. + n as f64) / input.K as f64)).round() as usize;
        }
    }
}

impl AnnealingState for State {
    #[inline]
    fn get_score(&self) -> usize {
        self.score
    }
    fn rollback(&mut self, diff: &DiffState) {
        *self = diff.state.clone();
    }
}

impl std::fmt::Display for State {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "{}", self.powers.iter().join(" "))?;
        write!(
            f,
            "{}",
            self.switches
                .iter()
                .map(|&b| if b { 1 } else { 0 })
                .join(" ")
        )
    }
}

struct DiffState {
    // TODO: その他復元に必要な情報を
    old_score: usize,
    state: State,
}

trait AnnealingTransition {
    fn transition(
        &self,
        state: &mut State,
        rng: &mut rand_pcg::Pcg64Mcg,
        input: &Input,
        dists_from_house_to_station: &Vec<Vec<usize>>,
        dists: &Vec<Vec<usize>>,
        path: &Vec<Vec<Vec<usize>>>,
    ) -> DiffState;
}

/// TODO: 名前変える。遷移毎に生成
struct OneRemoveTransition {}

impl AnnealingTransition for OneRemoveTransition {
    fn transition(
        &self,
        state: &mut State,
        rng: &mut rand_pcg::Pcg64Mcg,
        input: &Input,
        dists_from_house_to_station: &Vec<Vec<usize>>,
        dists: &Vec<Vec<usize>>,
        path: &Vec<Vec<Vec<usize>>>,
    ) -> DiffState {
        let old_score = state.get_score();
        let diff_state = DiffState {
            old_score,
            state: state.clone(),
        };
        let mut cand = vec![];
        for (i, used) in state.use_station.iter().enumerate() {
            if *used {
                cand.push(i);
            }
        }

        let remove_idx = cand[rng.gen_range(0, cand.len())];
        state.use_station[remove_idx] = false;

        choose_powers(&input, state, &dists_from_house_to_station, &dists, &path);

        state.compute_score(&input);
        diff_state
    }
}

struct AddAndRemoveTransition {}

impl AnnealingTransition for AddAndRemoveTransition {
    fn transition(
        &self,
        state: &mut State,
        rng: &mut rand_pcg::Pcg64Mcg,
        input: &Input,
        dists_from_house_to_station: &Vec<Vec<usize>>,
        dists: &Vec<Vec<usize>>,
        path: &Vec<Vec<Vec<usize>>>,
    ) -> DiffState {
        let old_score = state.get_score();
        let diff_state = DiffState {
            old_score,
            state: state.clone(),
        };
        let mut cand_remove = vec![];
        let mut cand_add = vec![];
        for (i, used) in state.use_station.iter().enumerate() {
            if *used {
                cand_remove.push(i);
            } else {
                cand_add.push(i);
            }
        }

        let remove_idx = cand_remove[rng.gen_range(0, cand_remove.len())];
        let add_idx = cand_add[rng.gen_range(0, cand_add.len())];
        state.use_station[remove_idx] = false;
        state.use_station[add_idx] = true;

        choose_powers(&input, state, &dists_from_house_to_station, &dists, &path);

        state.compute_score(&input);
        diff_state
    }
}

fn simulated_annealing<T: TemperatureScheduler>(
    state: &mut State,
    timekeeper: &TimeKeeper,
    temperature_scheduler: &T,
    input: &Input,
    dists_from_house_to_station: &Vec<Vec<usize>>,
    dists: &Vec<Vec<usize>>,
    path: &Vec<Vec<Vec<usize>>>,
) {
    let mut rng = rand_pcg::Pcg64Mcg::new(RNG_SEED);
    let mut best_state = state.clone();
    let mut temperature = temperature_scheduler.get_temperature(&timekeeper);
    let mut cnt = 0;
    let mut accepted_cnt = 0;
    let mut score_update_cnt = 0;
    let init_score = state.get_score();
    let update_per_cnt = 100;
    let neighbor_creater = OneRemoveTransition {};
    while timekeeper.intime() {
        cnt += 1;
        if cnt % update_per_cnt == 0 {
            temperature = temperature_scheduler.get_temperature(&timekeeper);
        }
        let old_score = state.get_score();
        let diff_state = neighbor_creater.transition(
            state,
            &mut rng,
            &input,
            &dists_from_house_to_station,
            &dists,
            &path,
        );
        let new_score = state.get_score();

        if new_score > old_score || rng.gen_bool(probability(new_score, old_score, temperature)) {
            if new_score > old_score {
                best_state = state.clone();
                score_update_cnt += 1;
            }
            accepted_cnt += 1
        } else {
            state.rollback(&diff_state);
        }
    }
    *state = best_state;
    eprintln!("-----------start Simulated Annealing--------------");
    eprintln!(
        "accepted / iter: {}/{}, {}%",
        accepted_cnt,
        cnt,
        (accepted_cnt * 100) as f64 / cnt as f64
    );
    eprintln!(
        "score update / iter: {}/{}, {}%",
        score_update_cnt,
        cnt,
        (score_update_cnt * 100) as f64 / cnt as f64
    );
    eprintln!("init score: {}", init_score);
    eprintln!("final score: {}", state.get_score());
    eprintln!("-----------end Simulated Annealing--------------");
}

#[derive(Debug, Clone)]
pub struct UnionFind {
    parent_or_size: Vec<i64>,
}

impl UnionFind {
    pub fn new(n: usize) -> Self {
        let parent_or_size = vec![-1; n];
        Self { parent_or_size }
    }

    pub fn root(&mut self, x: usize) -> usize {
        if self.parent_or_size[x] < 0 {
            x
        } else {
            self.parent_or_size[x] = self.root(self.parent_or_size[x] as usize) as i64;
            self.parent_or_size[x] as usize
        }
    }

    pub fn is_same(&mut self, x: usize, y: usize) -> bool {
        self.root(x) == self.root(y)
    }

    pub fn size(&mut self, x: usize) -> usize {
        let x = self.root(x);
        (-self.parent_or_size[x]) as usize
    }

    pub fn unite(&mut self, x: usize, y: usize) {
        let x = self.root(x);
        let y = self.root(y);
        if x == y {
            return;
        }
        // 大きい木の根に小さい木をくっつける
        if self.size(x) < self.size(y) {
            self.parent_or_size[y] += self.parent_or_size[x];
            self.parent_or_size[x] = y as i64;
        } else {
            self.parent_or_size[x] += self.parent_or_size[y];
            self.parent_or_size[y] = x as i64;
        }
    }
}

Submission Info

Compile Error

warning: method is never used: `new`
   --> src/main.rs:362:5
    |
362 |     fn new(initial_temperature: f64, end_temperature: f64) -> Self {
    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |
    = note: `#[warn(dead_code)]` on by default

warning: field is never read: `old_score`
   --> src/main.rs:473:5
    |
473 |     old_score: usize,
    |     ^^^^^^^^^^^^^^^^

Judge Result