#1195149 (Rust) No.5025 Forest Fire

• 問題: No.5025 Forest Fire
• ユーザー: ryomac
• 提出日時: 2026-10-24 17:24:17
• 言語: Rust (1.99.0 + proconio + num + itertools + ACL)
• 結果: AC
• 実行時間: 158 ms / 2,000 ms
• コード長: 19,400 bytes
• コンパイル時間: 979 ms
• コンパイル使用メモリ: 196,452 KB
• 実行使用メモリ: 9,924 KB
• スコア: 162,453,667
• 最終ジャッジ日時: 2026-10-24 17:31:24
• 合計ジャッジ時間: 63,917 ms
• ジャッジサーバーID （参考情報）: judge4_0 / <nil>

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

ソースコード

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

const FOREST_SIZE: usize = 50;
const CUT_LIMIT: usize = 100;
const DEFAULT_SAMPLES: usize = 50;
const CENTER_COUNT: usize = 150;
const CENTER_MIN: usize = 12;
const CENTER_SIDE: usize = 36;

#[derive(Clone, Copy)]
enum Direction {
    Up,
    Down,
    Left,
    Right,
}

impl Direction {
    fn is_horizontal(self) -> bool {
        matches!(self, Self::Left | Self::Right)
    }

    fn find_candidates(self, index: usize, n: usize) -> Vec<usize> {
        let row = index / n;
        let column = index % n;
        let (step, position) = match self {
            Self::Up if row > 0 => (row - 1, column),
            Self::Down if row + 1 < n => (row + 1, column),
            Self::Left if column > 0 => (column - 1, row),
            Self::Right if column + 1 < n => (column + 1, row),
            _ => return Vec::new(),
        };
        (position.saturating_sub(1)..=(position + 1).min(n - 1))
            .map(|candidate| get_index(step, candidate, n, self.is_horizontal()))
            .collect()
    }
}

fn get_index(step: usize, position: usize, n: usize, horizontal: bool) -> usize {
    if horizontal {
        position * n + step
    } else {
        step * n + position
    }
}

fn calculate_sum(grid: &[u8], cut: &[bool], index: usize, n: usize, horizontal: bool) -> u8 {
    let (step, position) = if horizontal {
        (index % n, index / n)
    } else {
        (index / n, index % n)
    };
    let find_neighbor = |direction: isize| {
        let mut neighbor = position as isize + direction;
        while (0..n as isize).contains(&neighbor) {
            let index = get_index(step, neighbor as usize, n, horizontal);
            // 伐採済みのマスを同じ向きに飛ばす。
            if cut[index] {
                neighbor += direction;
                continue;
            }
            return Some(grid[index]);
        }
        None
    };
    let (Some(left), Some(right)) = (find_neighbor(-1), find_neighbor(1)) else {
        return 0;
    };
    // 共通分母10を除き、乾燥度の合計を整数で比較する。
    left + right
}

fn select_next(
    grid: &[u8],
    cut: &[bool],
    index: usize,
    n: usize,
    direction: Direction,
) -> Option<usize> {
    let horizontal = direction.is_horizontal();
    let position = if horizontal { index / n } else { index % n };
    direction
        .find_candidates(index, n)
        .into_iter()
        .max_by_key(|&candidate| {
            let next_position = if horizontal {
                candidate / n
            } else {
                candidate % n
            };
            let sum = calculate_sum(grid, cut, candidate, n, horizontal);
            // 中央24.5との距離を2倍し、同点では直進を優先する。
            let center_distance = (2 * next_position).abs_diff(n - 1);
            (sum, n - center_distance, next_position == position)
        })
}

fn record_cut(grid: &[u8], cut: &mut [bool], cuts: &mut Vec<usize>, index: usize) {
    if grid[index] > 0 && !cut[index] {
        cut[index] = true;
        cuts.push(index);
    }
}

fn generate_cuts(grid: &[u8], center: usize, n: usize) -> Vec<usize> {
    let mut cut = vec![false; n * n];
    let mut cuts = Vec::with_capacity(2 * n);
    record_cut(grid, &mut cut, &mut cuts, center);
    for direction in [
        Direction::Up,
        Direction::Down,
        Direction::Left,
        Direction::Right,
    ] {
        let mut index = center;
        while let Some(next) = select_next(grid, &cut, index, n, direction) {
            index = next;
            record_cut(grid, &mut cut, &mut cuts, index);
        }
    }
    cuts
}

fn select_centers(n: usize) -> Vec<usize> {
    let mut candidates: Vec<usize> = (CENTER_MIN..CENTER_MIN + CENTER_SIDE)
        .flat_map(|row| (CENTER_MIN..CENTER_MIN + CENTER_SIDE).map(move |column| row * n + column))
        .collect();
    let mut rng = SplitMix64 { state: SEED };
    // 未選択の領域から一様に選び、先頭へ移す。
    for index in 0..CENTER_COUNT {
        let selected = index + rng.draw_index(candidates.len() - index);
        candidates.swap(index, selected);
    }
    candidates.truncate(CENTER_COUNT);
    candidates
}

fn select_solution(
    grid: &[u8],
    n: usize,
    centers: &[usize],
    evaluator: &mut Evaluator,
) -> (usize, Vec<usize>, u64) {
    let mut best_center = centers[0];
    let mut best_cuts = Vec::new();
    let mut best_sum = u64::MAX;
    for &center in centers {
        let mut cuts = generate_cuts(grid, center, n);
        let generated_count = cuts.len();
        cuts.truncate(CUT_LIMIT);
        if generated_count > CUT_LIMIT {
            eprintln!("cut_limit_exceeded generated_M={generated_count} output_M={CUT_LIMIT}");
        }
        let sum = evaluator.sum_component_squares(&cuts);
        // 共通の正の分母なので、整数の分子でEを比較する。
        if sum < best_sum {
            best_center = center;
            best_cuts = cuts;
            best_sum = sum;
        }
    }
    (best_center, best_cuts, best_sum)
}

fn main() {
    let started = Instant::now();
    let mut input = String::new();
    io::stdin().read_to_string(&mut input).expect("read stdin");
    let mut tokens = input.split_whitespace();
    let n: usize = tokens.next().expect("N").parse().expect("integer N");
    let k: usize = tokens.next().expect("K").parse().expect("integer K");
    assert_eq!(n, FOREST_SIZE, "N must be 50");
    assert_eq!(k, CUT_LIMIT, "K must be 100");
    let mut grid = Vec::with_capacity(n * n);
    for _ in 0..n {
        let row = tokens.next().expect("grid row");
        assert_eq!(row.len(), n, "grid row must contain N digits");
        for digit in row.bytes() {
            assert!(digit.is_ascii_digit(), "grid must contain digits");
            grid.push(digit - b'0');
        }
    }
    assert!(tokens.next().is_none(), "unexpected trailing input");
    let sample_count = match std::env::var("SAMPLES") {
        Ok(value) => value.parse::<usize>().expect("SAMPLES must be an integer"),
        Err(std::env::VarError::NotPresent) => DEFAULT_SAMPLES,
        Err(error) => panic!("invalid SAMPLES: {error}"),
    };
    let centers = select_centers(n);
    let mut evaluator = Evaluator::new(n, &grid, sample_count);
    let baseline = evaluator.sum_component_squares(&[]);
    let (center, cuts, final_sum) = select_solution(&grid, n, &centers, &mut evaluator);
    let mut output = io::BufWriter::new(io::stdout().lock());
    writeln!(output, "{}", cuts.len()).expect("write M");
    for &index in &cuts {
        writeln!(output, "{} {}", index / n, index % n).expect("write coordinates");
    }
    output.flush().expect("flush stdout");
    eprintln!(
        "elapsed_seconds={:.6} ci={} cj={} E={:.12} E0={:.12} M={}",
        started.elapsed().as_secs_f64(),
        center / n,
        center % n,
        evaluator.expectation(final_sum),
        evaluator.expectation(baseline),
        cuts.len(),
    );
}

const SEED: u64 = 42;
const DRYNESS_SCALE: u64 = 10;

struct SplitMix64 {
    state: u64,
}

impl SplitMix64 {
    fn draw_index(&mut self, count: usize) -> usize {
        let bound = count as u64;
        let threshold = bound.wrapping_neg() % bound;
        loop {
            let value = self.next();
            if value >= threshold {
                return (value % bound) as usize;
            }
        }
    }

    fn next(&mut self) -> u64 {
        self.state = self.state.wrapping_add(0x9e3779b97f4a7c15);
        let mut value = self.state;
        value = (value ^ (value >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
        value = (value ^ (value >> 27)).wrapping_mul(0x94d049bb133111eb);
        value ^ (value >> 31)
    }

    fn draw_dryness(&mut self) -> u64 {
        // 剰余の偏りを避け、0以上10未満を一様に生成する。
        let threshold = DRYNESS_SCALE.wrapping_neg() % DRYNESS_SCALE;
        loop {
            let value = self.next();
            if value >= threshold {
                return value % DRYNESS_SCALE;
            }
        }
    }
}

pub struct Evaluator {
    n: usize,
    stride: usize,
    samples: Vec<Vec<u8>>,
    remaining: Vec<u8>,
    stack: Vec<usize>,
}

impl Evaluator {
    pub fn new(n: usize, grid: &[u8], sample_count: usize) -> Self {
        assert!(sample_count > 0, "SAMPLES must be positive");
        assert_eq!(grid.len(), n * n);
        let stride = n + 2;
        let area = stride * stride;
        let mut rng = SplitMix64 { state: SEED };
        let samples = (0..sample_count)
            .map(|_| {
                // 外周を木なしにして、4近傍の境界判定を不要にする。
                let mut sample = vec![0; area];
                for (index, &dryness) in grid.iter().enumerate() {
                    if dryness > 0 && rng.draw_dryness() < u64::from(dryness) {
                        sample[(index / n + 1) * stride + index % n + 1] = 1;
                    }
                }
                sample
            })
            .collect();
        Self {
            n,
            stride,
            samples,
            remaining: vec![0; area],
            stack: Vec::with_capacity(n * n),
        }
    }

    pub fn sum_component_squares(&mut self, cuts: &[usize]) -> u64 {
        let mut total = 0;
        for sample in &self.samples {
            self.remaining.copy_from_slice(sample);
            for &index in cuts {
                self.remaining[(index / self.n + 1) * self.stride + index % self.n + 1] = 0;
            }
            for row in 1..=self.n {
                for column in 1..=self.n {
                    let start = row * self.stride + column;
                    if self.remaining[start] == 0 {
                        continue;
                    }
                    self.remaining[start] = 0;
                    self.stack.push(start);
                    let mut size = 0_u64;
                    while let Some(index) = self.stack.pop() {
                        size += 1;
                        for neighbor in [
                            index - self.stride,
                            index + self.stride,
                            index - 1,
                            index + 1,
                        ] {
                            if self.remaining[neighbor] != 0 {
                                self.remaining[neighbor] = 0;
                                self.stack.push(neighbor);
                            }
                        }
                    }
                    total += size * size;
                }
            }
        }
        total
    }

    pub fn expectation(&self, sum: u64) -> f64 {
        sum as f64 / (self.n * self.n) as f64 / self.samples.len() as f64
    }
}

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

    #[test]
    fn select_thirty_unique_reproducible_centers_in_central_region() {
        let centers = select_centers(FOREST_SIZE);
        assert_eq!(centers.len(), 30);
        assert_eq!(centers, select_centers(FOREST_SIZE));
        for &center in &centers {
            assert!((12..=36).contains(&(center / FOREST_SIZE)));
            assert!((12..=36).contains(&(center % FOREST_SIZE)));
        }
        let mut unique = centers.clone();
        unique.sort_unstable();
        unique.dedup();
        assert_eq!(unique.len(), 30);
    }

    #[test]
    fn select_minimum_e_with_common_samples_and_keep_first_tie() {
        let centers = select_centers(FOREST_SIZE);
        for grid in [
            vec![0; FOREST_SIZE * FOREST_SIZE],
            vec![9; FOREST_SIZE * FOREST_SIZE],
        ] {
            let mut evaluator = Evaluator::new(FOREST_SIZE, &grid, DEFAULT_SAMPLES);
            let original_samples = evaluator.samples.clone();
            let expected: Vec<_> = centers
                .iter()
                .map(|&center| {
                    let mut cuts = generate_cuts(&grid, center, FOREST_SIZE);
                    cuts.truncate(CUT_LIMIT);
                    evaluator.sum_component_squares(&cuts)
                })
                .collect();
            let minimum = *expected.iter().min().unwrap();
            let first_minimum = expected.iter().position(|&sum| sum == minimum).unwrap();
            let (center, cuts, sum) = select_solution(&grid, FOREST_SIZE, &centers, &mut evaluator);
            assert_eq!(center, centers[first_minimum]);
            assert_eq!(sum, minimum);
            assert_eq!(sum, evaluator.sum_component_squares(&cuts));
            assert_eq!(evaluator.samples, original_samples);
        }
    }

    #[test]
    fn find_candidates_in_all_four_directions() {
        for (direction, expected) in [
            (Direction::Up, vec![6, 7, 8]),
            (Direction::Down, vec![16, 17, 18]),
            (Direction::Left, vec![6, 11, 16]),
            (Direction::Right, vec![8, 13, 18]),
        ] {
            assert_eq!(direction.find_candidates(12, 5), expected);
        }
        assert_eq!(Direction::Up.find_candidates(5, 5), vec![0, 1]);
        assert_eq!(Direction::Down.find_candidates(19, 5), vec![23, 24]);
        assert_eq!(Direction::Left.find_candidates(1, 5), vec![0, 5]);
        assert_eq!(Direction::Right.find_candidates(21, 5), vec![17, 22]);
        for (direction, index) in [
            (Direction::Up, 2),
            (Direction::Down, 22),
            (Direction::Left, 10),
            (Direction::Right, 14),
        ] {
            assert!(direction.find_candidates(index, 5).is_empty());
        }
    }

    #[test]
    fn calculate_values_at_edges_and_zero_dryness() {
        for horizontal in [false, true] {
            let mut grid = vec![0; 25];
            for (position, dryness) in [9, 1, 4, 0, 8].into_iter().enumerate() {
                grid[get_index(2, position, 5, horizontal)] = dryness;
            }
            let cut = vec![false; 25];
            for (position, expected) in [0, 13, 1, 12, 0].into_iter().enumerate() {
                assert_eq!(
                    calculate_sum(
                        &grid,
                        &cut,
                        get_index(2, position, 5, horizontal),
                        5,
                        horizontal
                    ),
                    expected
                );
            }
        }
    }

    #[test]
    fn skip_cuts_in_both_axes_and_stop_at_zero_or_edge() {
        for horizontal in [false, true] {
            let mut grid = vec![1; 49];
            let mut cut = vec![false; 49];
            grid[get_index(3, 0, 7, horizontal)] = 9;
            grid[get_index(3, 6, 7, horizontal)] = 4;
            for position in [1, 2, 4, 5] {
                cut[get_index(3, position, 7, horizontal)] = true;
            }
            assert_eq!(calculate_sum(&grid, &cut, 24, 7, horizontal), 13);
            grid[get_index(3, 2, 7, horizontal)] = 0;
            cut[get_index(3, 2, 7, horizontal)] = false;
            assert_eq!(calculate_sum(&grid, &cut, 24, 7, horizontal), 4);
            cut[get_index(3, 2, 7, horizontal)] = true;
            cut[get_index(3, 0, 7, horizontal)] = true;
            assert_eq!(calculate_sum(&grid, &cut, 24, 7, horizontal), 0);
            cut[get_index(3, 0, 7, horizontal)] = false;
            cut[get_index(3, 6, 7, horizontal)] = true;
            assert_eq!(calculate_sum(&grid, &cut, 24, 7, horizontal), 0);
        }
    }

    #[test]
    fn select_largest_value_before_center() {
        for direction in [
            Direction::Up,
            Direction::Down,
            Direction::Left,
            Direction::Right,
        ] {
            let horizontal = direction.is_horizontal();
            let next_step = match direction {
                Direction::Up | Direction::Left => 1,
                Direction::Down | Direction::Right => 3,
            };
            let mut grid = vec![0; 25];
            for (position, dryness) in [9, 1, 8, 9, 0].into_iter().enumerate() {
                grid[get_index(next_step, position, 5, horizontal)] = dryness;
            }
            assert_eq!(
                select_next(&grid, &[false; 25], 12, 5, direction),
                Some(get_index(next_step, 1, 5, horizontal))
            );
        }
    }

    #[test]
    fn break_ties_by_center_then_straight_in_all_directions() {
        let grid = vec![1; 2500];
        let cut = vec![false; 2500];
        for direction in [
            Direction::Up,
            Direction::Down,
            Direction::Left,
            Direction::Right,
        ] {
            let horizontal = direction.is_horizontal();
            let next_step = match direction {
                Direction::Up | Direction::Left => 9,
                Direction::Down | Direction::Right => 11,
            };
            for (position, expected) in [(23, 24), (24, 24), (25, 25), (26, 25)] {
                assert_eq!(
                    select_next(
                        &grid,
                        &cut,
                        get_index(10, position, 50, horizontal),
                        50,
                        direction
                    ),
                    Some(get_index(next_step, expected, 50, horizontal))
                );
            }
        }
    }

    #[test]
    fn advance_into_cut_or_zero_candidates_without_recording_them() {
        let grid = vec![0; 25];
        let mut cut = vec![false; 25];
        cut[7] = true;
        assert_eq!(select_next(&grid, &cut, 12, 5, Direction::Up), Some(7));
        let mut cuts = vec![7];
        record_cut(&grid, &mut cut, &mut cuts, 7);
        record_cut(&grid, &mut cut, &mut cuts, 12);
        assert_eq!(cuts, vec![7]);
        let mut grid = vec![0; 25];
        grid[1..4].fill(9);
        // 中心と途中のd=0を通過して端の木に到達する。
        assert_eq!(generate_cuts(&grid, 12, 5), vec![2]);
        assert!(generate_cuts(&[0; 25], 12, 5).is_empty());
    }

    #[test]
    fn record_center_then_up_down_left_right_without_duplicates() {
        assert_eq!(
            generate_cuts(&[1; 25], 12, 5),
            vec![12, 7, 2, 17, 22, 11, 10, 13, 14]
        );
        let grid = vec![9; 2500];
        let cuts = generate_cuts(&grid, 24 * 50 + 24, 50);
        assert!(cuts.len() <= CUT_LIMIT);
        let mut unique = cuts.clone();
        unique.sort_unstable();
        unique.dedup();
        assert_eq!(unique.len(), cuts.len());
    }

    #[test]
    fn evaluate_components_and_cuts() {
        let mut evaluator = Evaluator::new(3, &[0; 9], 2);
        for index in [6, 7, 8, 13, 16, 17] {
            evaluator.samples[0][index] = 1;
        }
        evaluator.samples[1][6] = 1;
        assert_eq!(evaluator.sum_component_squares(&[]), 21);
        assert_eq!(evaluator.sum_component_squares(&[1, 3, 4, 5, 7]), 4);
        assert_eq!(evaluator.expectation(21), 21.0 / 18.0);
        assert_eq!(evaluator.sum_component_squares(&[]), 21);
    }

    #[test]
    fn generate_reproducible_samples_without_water() {
        let grid = [0, 9, 1, 0];
        let first = Evaluator::new(2, &grid, 50);
        let second = Evaluator::new(2, &grid, 50);
        assert_eq!(first.samples, second.samples);
        for sample in &first.samples {
            assert_eq!(sample[5], 0);
            assert_eq!(sample[10], 0);
        }
    }
}