#1195075 (Rust) No.5025 Forest Fire

• 問題: No.5025 Forest Fire
• ユーザー: ryomac
• 提出日時: 2026-10-24 16:41:21
• 言語: Rust (1.99.0 + proconio + num + itertools + ACL)
• 結果: AC
• 実行時間: 95 ms / 2,000 ms
• コード長: 19,404 bytes
• コンパイル時間: 865 ms
• コンパイル使用メモリ: 182,824 KB
• 実行使用メモリ: 9,928 KB
• スコア: 161,089,120
• 最終ジャッジ日時: 2026-10-24 16:42:27
• 合計ジャッジ時間: 62,348 ms
• ジャッジサーバーID （参考情報）: judge2_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;

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],
    step: usize,
    position: usize,
    n: usize,
    horizontal: bool,
) -> u8 {
    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);
            // 2本目では1本目の伐採済みマスを同じ向きに飛ばす。
            if horizontal && 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],
    step: usize,
    position: usize,
    n: usize,
    horizontal: bool,
) -> Option<usize> {
    let mut best = None;
    let mut best_key = None;
    for candidate in position.saturating_sub(1)..=(position + 1).min(n - 1) {
        if cut[get_index(step, candidate, n, horizontal)] {
            continue;
        }
        let sum = calculate_sum(grid, cut, step, candidate, n, horizontal);
        // 中央との距離を2倍して、24.5との比較も整数で行う。
        let center_distance = (2 * candidate).abs_diff(n - 1);
        let key = (sum, n - center_distance, candidate == position);
        if best_key.is_none_or(|previous| key > previous) {
            best = Some(candidate);
            best_key = Some(key);
        }
    }
    best
}

fn generate_path(
    grid: &[u8],
    base_cuts: &[usize],
    start: usize,
    n: usize,
    horizontal: bool,
) -> Vec<usize> {
    let mut cut = vec![false; n * n];
    for &index in base_cuts {
        cut[index] = true;
    }
    let mut path = Vec::with_capacity(n);
    let mut position = start;
    for step in 0..n {
        if step > 0 {
            let Some(next) = select_next(grid, &cut, step, position, n, horizontal) else {
                continue;
            };
            position = next;
        }
        let index = get_index(step, position, n, horizontal);
        if grid[index] > 0 && !cut[index] {
            cut[index] = true;
            path.push(index);
        }
    }
    path
}

fn select_path(
    evaluator: &mut Evaluator,
    grid: &[u8],
    base_cuts: &[usize],
    horizontal: bool,
) -> (usize, Vec<usize>, u64) {
    let mut best_start = 0;
    let mut best_path = Vec::new();
    let mut best_sum = u64::MAX;
    for start in 0..evaluator.n {
        let path = generate_path(grid, base_cuts, start, evaluator.n, horizontal);
        let mut cuts = base_cuts.to_vec();
        cuts.extend_from_slice(&path);
        let sum = evaluator.sum_component_squares(&cuts);
        // 分母は全候補で共通。同値では開始位置の小さい経路を維持する。
        if sum < best_sum {
            best_start = start;
            best_path = path;
            best_sum = sum;
        }
    }
    (best_start, best_path, 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 mut evaluator = Evaluator::new(n, &grid, sample_count);
    let baseline = evaluator.sum_component_squares(&[]);
    let (column, mut cuts, _) = select_path(&mut evaluator, &grid, &[], false);
    let (row, second_path, best_sum) = select_path(&mut evaluator, &grid, &cuts, true);
    cuts.extend_from_slice(&second_path);
    assert!(cuts.len() <= k);
    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} c={} r={} E={:.12} E0={:.12} M={}",
        started.elapsed().as_secs_f64(),
        column,
        row,
        evaluator.expectation(best_sum),
        evaluator.expectation(baseline),
        cuts.len(),
    );
}

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

struct SplitMix64 {
    state: u64,
}

impl SplitMix64 {
    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 retain_first_path_on_ties() {
        let grid = [0; 9];
        let mut evaluator = Evaluator::new(3, &grid, 2);
        assert_eq!(
            select_path(&mut evaluator, &grid, &[], false),
            (0, vec![], 0)
        );
        assert_eq!(
            select_path(&mut evaluator, &grid, &[], true),
            (0, vec![], 0)
        );
    }

    #[test]
    fn evaluate_components_and_cuts() {
        let mut evaluator = Evaluator::new(3, &[0; 9], 2);
        evaluator.samples[0] = vec![0; 25];
        evaluator.samples[1] = vec![0; 25];
        for index in [6, 7, 8, 13, 16, 17] {
            evaluator.samples[0][index] = 1;
        }
        evaluator.samples[1][6] = 1;
        // 第1サンプルは大きさ4と2、第2サンプルは大きさ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 select_smallest_expectation_with_base_cuts() {
        let grid = [9; 9];
        for horizontal in [false, true] {
            let mut evaluator = Evaluator::new(3, &grid, 1);
            evaluator.samples[0].fill(0);
            for position in 0..3 {
                let index = get_index(0, position, 3, horizontal);
                evaluator.samples[0][(index / 3 + 1) * 5 + index % 3 + 1] = 1;
            }
            let expected = generate_path(&grid, &[], 1, 3, horizontal);
            assert_eq!(
                select_path(&mut evaluator, &grid, &[], horizontal),
                (1, expected, 2)
            );
            let base_cuts = [get_index(0, 1, 3, horizontal)];
            let expected = generate_path(&grid, &base_cuts, 0, 3, horizontal);
            assert_eq!(
                select_path(&mut evaluator, &grid, &base_cuts, horizontal),
                (0, expected, 1)
            );
        }
    }

    #[test]
    fn generate_same_paths_with_base_cuts() {
        let grid = [9, 1, 4, 0, 7, 0, 2, 6, 8];
        let transposed: Vec<_> = (0..9).map(|i| grid[i % 3 * 3 + i / 3]).collect();
        for (board, horizontal, neighbor) in [(&grid[..], false, 0), (&transposed[..], true, 0)] {
            let expected: Vec<_> = (0..3)
                .map(|step| get_index(step, 1, 3, horizontal))
                .collect();
            assert_eq!(generate_path(board, &[], 1, 3, horizontal), expected);
            assert_eq!(
                generate_path(board, &[neighbor], 1, 3, horizontal),
                expected
            );
            assert_eq!(
                generate_path(board, &[expected[0]], 1, 3, horizontal),
                expected[1..].to_vec()
            );
        }
    }

    #[test]
    fn return_zero_when_either_neighbor_is_outside() {
        let grid = [9, 1, 4, 0, 7, 0, 0, 0, 0];
        let cut = vec![false; 9];
        let transposed: Vec<_> = (0..9).map(|i| grid[i % 3 * 3 + i / 3]).collect();
        for (board, horizontal) in [(&grid[..], false), (&transposed[..], true)] {
            assert_eq!(calculate_sum(board, &cut, 0, 0, 3, horizontal), 0);
            assert_eq!(calculate_sum(board, &cut, 0, 2, 3, horizontal), 0);
            assert_eq!(calculate_sum(board, &cut, 0, 1, 3, horizontal), 13);
        }
    }

    #[test]
    fn use_zero_dryness_as_zero_probability() {
        let grid = [0, 1, 4, 9, 7, 0, 0, 1, 0];
        let cut = vec![false; 9];
        let transposed: Vec<_> = (0..9).map(|i| grid[i % 3 * 3 + i / 3]).collect();
        for (board, horizontal) in [(&grid[..], false), (&transposed[..], true)] {
            assert_eq!(calculate_sum(board, &cut, 0, 1, 3, horizontal), 4);
            assert_eq!(calculate_sum(board, &cut, 1, 1, 3, horizontal), 9);
            assert_eq!(calculate_sum(board, &cut, 2, 1, 3, horizontal), 0);
        }
    }

    #[test]
    fn skip_consecutive_first_path_cuts_in_both_directions() {
        let mut grid = vec![1; 49];
        grid[3] = 9;
        grid[45] = 4;
        let mut cut = vec![false; 49];
        for index in [10, 17, 31, 38] {
            cut[index] = true;
        }
        assert_eq!(calculate_sum(&grid, &cut, 3, 3, 7, true), 13);

        // 伐採済みでないd=0のマスでは探索を止める。
        grid[10] = 0;
        cut[10] = false;
        assert_eq!(calculate_sum(&grid, &cut, 3, 3, 7, true), 4);
    }

    #[test]
    fn return_zero_when_skipping_cuts_reaches_either_edge() {
        let grid = vec![9; 25];
        for indices in [[2, 7], [17, 22]] {
            let mut cut = vec![false; 25];
            for index in indices {
                cut[index] = true;
            }
            assert_eq!(calculate_sum(&grid, &cut, 2, 2, 5, true), 0);
        }
    }

    #[test]
    fn keep_first_path_comparison_on_immediate_neighbors() {
        let grid = [9, 1, 4, 0, 7, 0, 0, 0, 0];
        let mut cut = vec![false; 9];
        cut[0] = true;
        assert_eq!(calculate_sum(&grid, &cut, 0, 1, 3, false), 13);
    }

    #[test]
    fn select_largest_sum_before_center() {
        for horizontal in [false, true] {
            let mut grid = vec![0; 25];
            for (position, dryness) in [9, 1, 8, 9, 0].into_iter().enumerate() {
                grid[get_index(1, position, 5, horizontal)] = dryness;
            }
            let cut = vec![false; 25];
            assert_eq!(calculate_sum(&grid, &cut, 1, 1, 5, horizontal), 17);
            assert_eq!(calculate_sum(&grid, &cut, 1, 2, 5, horizontal), 10);
            assert_eq!(calculate_sum(&grid, &cut, 1, 3, 5, horizontal), 8);
            assert_eq!(select_next(&grid, &cut, 1, 2, 5, horizontal), Some(1));
        }
    }

    #[test]
    fn break_ties_by_center_then_straight() {
        let grid = vec![1; FOREST_SIZE * FOREST_SIZE];
        let cut = vec![false; grid.len()];
        for horizontal in [false, true] {
            assert_eq!(select_next(&grid, &cut, 1, 23, 50, horizontal), Some(24));
            assert_eq!(select_next(&grid, &cut, 1, 24, 50, horizontal), Some(24));
            assert_eq!(select_next(&grid, &cut, 1, 25, 50, horizontal), Some(25));
            assert_eq!(select_next(&grid, &cut, 1, 26, 50, horizontal), Some(25));
        }
    }

    #[test]
    fn advance_through_zero_dryness_without_cutting_it() {
        let mut grid = vec![0; 25];
        grid[5..10].copy_from_slice(&[9, 0, 0, 0, 0]);
        grid[10..15].copy_from_slice(&[0, 0, 7, 9, 0]);
        let transposed: Vec<_> = (0..25).map(|i| grid[i % 5 * 5 + i / 5]).collect();
        for (board, horizontal) in [(&grid[..], false), (&transposed[..], true)] {
            let cut = vec![false; 25];
            assert_eq!(select_next(board, &cut, 1, 0, 5, horizontal), Some(1));
            assert_eq!(select_next(board, &cut, 2, 1, 5, horizontal), Some(2));
            // d=0の開始位置と移動先は出力せず、その移動先から次を選ぶ。
            assert_eq!(generate_path(board, &[], 0, 5, horizontal), vec![12]);
        }
    }

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

    #[test]
    fn exclude_cut_candidates_but_keep_zero_dryness_candidates() {
        let mut grid = vec![0; 25];
        grid[8] = 9;
        grid[12] = 7;
        grid[22] = 6;
        let mut cut = vec![false; 25];
        cut[8] = true;
        assert_eq!(select_next(&grid, &cut, 1, 2, 5, false), Some(2));
        assert_eq!(generate_path(&grid, &[8], 2, 5, false), vec![]);
        let transposed: Vec<_> = (0..25).map(|i| grid[i % 5 * 5 + i / 5]).collect();
        assert_eq!(generate_path(&transposed, &[16], 2, 5, true), vec![]);
        for horizontal in [false, true] {
            let board = vec![1; 9];
            let mut cut = vec![false; 9];
            cut[get_index(1, 1, 3, horizontal)] = true;
            assert_eq!(select_next(&board, &cut, 1, 1, 3, horizontal), Some(0));
            cut[get_index(1, 0, 3, horizontal)] = true;
            cut[get_index(1, 2, 3, horizontal)] = true;
            assert_eq!(select_next(&board, &cut, 1, 1, 3, horizontal), None);
        }
    }

    #[test]
    fn exclude_cut_start_and_cut_only_existing_trees() {
        let grid = vec![1; 9];
        assert_eq!(generate_path(&grid, &[1], 1, 3, false), vec![4, 7]);
        assert_eq!(generate_path(&grid, &[3], 1, 3, true), vec![4, 5]);

        let mut evaluator = Evaluator::new(3, &grid, 2);
        let (_, mut cuts, _) = select_path(&mut evaluator, &grid, &[], false);
        let (_, second, sum) = select_path(&mut evaluator, &grid, &cuts, true);
        cuts.extend(second);
        let mut unique = cuts.clone();
        unique.sort_unstable();
        unique.dedup();
        assert_eq!(unique.len(), cuts.len());
        assert!(cuts.len() <= 6);
        assert_eq!(sum, evaluator.sum_component_squares(&cuts));
    }

    #[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);
        }
    }
}