#1195014 (Rust) No.5025 Forest Fire

• 問題: No.5025 Forest Fire
• ユーザー: ryomac
• 提出日時: 2026-10-24 15:53:52
• 言語: Rust (1.99.0 + proconio + num + itertools + ACL)
• 結果: TLE
• 実行時間: -
• コード長: 7,698 bytes
• コンパイル時間: 5,051 ms
• コンパイル使用メモリ: 190,528 KB
• 実行使用メモリ: 9,928 KB
• スコア: 0
• 最終ジャッジ日時: 2026-10-24 15:56:55
• 合計ジャッジ時間: 182,753 ms
• ジャッジサーバーID （参考情報）: judge4_0 / <nil>

ファイルパターン	結果
other	TLE * 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 select_cross(evaluator: &mut Evaluator, n: usize) -> ((usize, usize), u64) {
    let mut best_cross = (0, 0);
    let mut best_sum = u64::MAX;
    for row in 0..n {
        for column in 0..n {
            let sum = evaluator.sum_component_squares(Some((row, column)));
            // 分母は全候補で共通。同値では先に見つかった組を維持する。
            if sum < best_sum {
                best_sum = sum;
                best_cross = (row, column);
            }
        }
    }
    (best_cross, 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(None);
    let ((row, column), best_sum) = select_cross(&mut evaluator, n);
    let cuts: Vec<_> = grid
        .iter()
        .enumerate()
        .filter(|&(index, &dryness)| dryness > 0 && (index / n == row || index % n == column))
        .map(|(index, _)| (index / n, index % n))
        .collect();
    assert!(cuts.len() <= k);
    let mut output = io::BufWriter::new(io::stdout().lock());
    writeln!(output, "{}", cuts.len()).expect("write M");
    for (row, column) in cuts {
        writeln!(output, "{row} {column}").expect("write coordinates");
    }
    output.flush().expect("flush stdout");
    eprintln!(
        "elapsed_seconds={:.6} r={} c={} E={:.12} E0={:.12}",
        started.elapsed().as_secs_f64(),
        row,
        column,
        evaluator.expectation(best_sum),
        evaluator.expectation(baseline),
    );
}

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, cross: Option<(usize, usize)>) -> u64 {
        let mut total = 0;
        for sample in &self.samples {
            self.remaining.copy_from_slice(sample);
            if let Some((row, column)) = cross {
                for position in 1..=self.n {
                    self.remaining[(row + 1) * self.stride + position] = 0;
                    self.remaining[position * self.stride + column + 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_cross_on_ties() {
        let mut evaluator = Evaluator::new(3, &[0; 9], 2);
        assert_eq!(select_cross(&mut evaluator, 3), ((0, 0), 0));
    }

    #[test]
    fn evaluate_components_and_cross() {
        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(None), 21);
        assert_eq!(evaluator.sum_component_squares(Some((1, 1))), 4);
        assert_eq!(evaluator.expectation(21), 21.0 / 18.0);
        assert_eq!(evaluator.sum_component_squares(None), 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);
        }
    }
}