#1195166 (Rust) No.5025 Forest Fire

• 問題: No.5025 Forest Fire
• ユーザー: g4np0n_kyopro
• 提出日時: 2026-10-24 17:29:58
• 言語: Rust (1.99.0 + proconio + num + itertools + ACL)
• 結果: TLE
• 実行時間: -
• コード長: 14,817 bytes
• コンパイル時間: 2,067 ms
• コンパイル使用メモリ: 203,596 KB
• 実行使用メモリ: 9,940 KB
• スコア: 168,329,020
• 最終ジャッジ日時: 2026-10-24 17:45:34
• 合計ジャッジ時間: 153,091 ms
• ジャッジサーバーID （参考情報）: judge3_0 / <nil>

ファイルパターン	結果
other	AC * 47 TLE * 3

ソースコード

use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hasher};
use std::io::{self, Read};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::OnceLock;

// 定数定義
const EXPECTATION_SAMPLES: usize = 1000;
const HIGH_RISK_PERCENTS: [usize; 9] = [20, 30, 40, 50, 60, 70, 80, 90, 100];
const NONE: usize = usize::MAX;

struct BurnEstimate {
    expected_burned: f64,
    /// 外周パディングを除いたN×N配列。添字は0始まりの盤面座標。
    /// 木のないマス・伐採済みのマスは0。
    burn_probabilities: Vec<Vec<f64>>,
}

fn next_expectation_seed() -> u64 {
    // 起動ごとにランダムな初期値を用意し、呼び出しごとにシードを進める。
    static STATE: OnceLock<AtomicU64> = OnceLock::new();
    let state = STATE.get_or_init(|| AtomicU64::new(RandomState::new().build_hasher().finish()));
    let mut seed = state.fetch_add(0x9E3779B97F4A7C15, Ordering::Relaxed);
    splitmix64(&mut seed)
}

fn splitmix64(state: &mut u64) -> u64 {
    *state = state.wrapping_add(0x9E3779B97F4A7C15);
    let mut z = *state;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
    z ^ (z >> 31)
}

// 1行を64ビットにまとめ、横方向の展開と隣接行への伝播をまとめて処理する。
// visitedの管理だけをビット化し、連結成分と各マスへの加算値は通常の探索と同じ。
struct BitComponents {
    remaining: Vec<u64>,
    pending: Vec<u64>,
    component: Vec<u64>,
    queue: Vec<usize>,
    touched_rows: Vec<usize>,
}

impl BitComponents {
    fn new(n: usize) -> Self {
        Self {
            remaining: vec![0; n],
            pending: vec![0; n],
            component: vec![0; n],
            queue: Vec::with_capacity(n * n),
            touched_rows: Vec::with_capacity(n),
        }
    }

    #[inline]
    fn extend_row(&mut self, row: usize, bits: u64) -> u64 {
        let fresh = self.remaining[row] & bits;
        if fresh == 0 {
            return 0;
        }
        self.remaining[row] &= !fresh;
        if self.component[row] == 0 {
            self.touched_rows.push(row);
        }
        self.component[row] |= fresh;
        if self.pending[row] == 0 {
            self.queue.push(row);
        }
        self.pending[row] |= fresh;
        u64::from(fresh.count_ones())
    }

    fn accumulate(&mut self, width: usize, burn_weights: &mut [u64]) -> u64 {
        let n = self.remaining.len();
        let mut sum = 0;
        for start_row in 0..n {
            while self.remaining[start_row] != 0 {
                self.queue.clear();
                self.touched_rows.clear();
                let start = self.remaining[start_row] & self.remaining[start_row].wrapping_neg();
                let mut size = self.extend_row(start_row, start);
                let mut cursor = 0;
                while cursor < self.queue.len() {
                    let row = self.queue[cursor];
                    cursor += 1;
                    let mut bits = self.pending[row];
                    self.pending[row] = 0;

                    // この行で到達できる左右の木をまとめて取り込む。
                    loop {
                        let fresh = ((bits << 1) | (bits >> 1)) & self.remaining[row];
                        if fresh == 0 {
                            break;
                        }
                        self.remaining[row] &= !fresh;
                        self.component[row] |= fresh;
                        bits |= fresh;
                        size += u64::from(fresh.count_ones());
                    }
                    if row > 0 {
                        size += self.extend_row(row - 1, bits);
                    }
                    if row + 1 < n {
                        size += self.extend_row(row + 1, bits);
                    }
                }
                for &row in &self.touched_rows {
                    let mut bits = self.component[row];
                    let offset = (row + 1) * width + 1;
                    while bits != 0 {
                        let col = bits.trailing_zeros() as usize;
                        burn_weights[offset + col] += size;
                        bits &= bits - 1;
                    }
                    self.component[row] = 0;
                }
                sum += size * size;
            }
        }
        sum
    }
}

// 各サンプルで、連結成分サイズ/N²を所属マスの焼失確率に加算する。
fn estimate_expected_burned(dryness: &[Vec<u32>], sample_count: usize) -> BurnEstimate {
    assert!(sample_count > 0, "サンプル数は1以上が必要です");
    let width = dryness.len();
    assert!(width >= 3, "盤面には1マス以上と外周が必要です");
    assert!(dryness.iter().all(|row| row.len() == width));
    let n = width - 2;
    assert!(n <= 64);
    let mut thresholds = vec![0_u64; width * width];
    for i in 0..width {
        for j in 0..width {
            let d = dryness[i][j];
            assert!(d <= 9);
            if i == 0 || j == 0 || i == width - 1 || j == width - 1 {
                assert_eq!(d, 0, "外周は0で埋めてください");
            }
            // 配布ツールと同じ32ビット乱数による乾燥判定。
            thresholds[i * width + j] = (u64::from(d) << 32) / 10;
        }
    }

    let mut state = next_expectation_seed();
    let mut burn_weights = vec![0_u64; width * width];
    let mut total = 0.0;
    let mut components = BitComponents::new(n);
    for _ in 0..sample_count {
        for row in 0..n {
            let mut bits = 0_u64;
            let offset = (row + 1) * width + 1;
            for col in 0..n {
                let random = splitmix64(&mut state) as u32;
                bits |= u64::from(u64::from(random) < thresholds[offset + col]) << col;
            }
            components.remaining[row] = bits;
        }
        total += components.accumulate(width, &mut burn_weights) as f64;
    }
    let denominator = sample_count as f64 * (n * n) as f64;
    let burn_probabilities = (1..=n)
        .map(|i| {
            (1..=n)
                .map(|j| burn_weights[i * width + j] as f64 / denominator)
                .collect()
        })
        .collect();
    BurnEstimate {
        expected_burned: total / denominator,
        burn_probabilities,
    }
}

struct Board {
    n: usize,
    dryness: Vec<u32>,
    neighbors: Vec<[usize; 4]>,
}

impl Board {
    fn new(dryness: &[Vec<u32>]) -> Self {
        let n = dryness.len() - 2;
        let mut values = Vec::with_capacity(n * n);
        let mut neighbors = Vec::with_capacity(n * n);
        for i in 0..n {
            for j in 0..n {
                values.push(dryness[i + 1][j + 1]);
                neighbors.push([
                    if i > 0 { (i - 1) * n + j } else { NONE },
                    if i + 1 < n { (i + 1) * n + j } else { NONE },
                    if j > 0 { i * n + j - 1 } else { NONE },
                    if j + 1 < n { i * n + j + 1 } else { NONE },
                ]);
            }
        }
        Self {
            n,
            dryness: values,
            neighbors,
        }
    }
}

#[derive(Clone, Copy)]
enum CutDirection {
    X,
    Y,
    XPlusY,
    XMinusY,
}

impl CutDirection {
    fn project(self, cell: usize, n: usize) -> i32 {
        let x = (cell % n) as i32;
        let y = (cell / n) as i32;
        match self {
            Self::X => x,
            Self::Y => y,
            Self::XPlusY => x + y,
            Self::XMinusY => x - y,
        }
    }
}

struct BalancedCut {
    cells: Vec<usize>,
    imbalance: usize,
}

/// 各方向でaを全探索し、左右の要素数の差が最小の切断を求める。
/// その4候補の中で、線上の木の本数が最小のものを返す。
/// 切断線上の木は左右のどちらにも含めず、両側に木が残る切断のみ扱う。
/// 同じ不均衡なら短い切断、同じ長さなら不均衡が小さい切断を優先する。
fn find_balanced_cut(cells: &[usize], n: usize) -> Option<BalancedCut> {
    let mut best: Option<BalancedCut> = None;
    for direction in [
        CutDirection::X,
        CutDirection::Y,
        CutDirection::XPlusY,
        CutDirection::XMinusY,
    ] {
        let values: Vec<i32> = cells.iter().map(|&v| direction.project(v, n)).collect();
        let Some(&minimum) = values.iter().min() else {
            return None;
        };
        let maximum = *values.iter().max().unwrap();
        let mut direction_best: Option<BalancedCut> = None;
        for a in minimum..=maximum {
            let left_count = values.iter().filter(|&&v| v < a).count();
            let right_count = values.iter().filter(|&&v| v > a).count();
            if left_count == 0 || right_count == 0 {
                continue;
            }
            let cut_cells: Vec<usize> = cells
                .iter()
                .copied()
                .zip(&values)
                .filter_map(|(v, &value)| (value == a).then_some(v))
                .collect();
            if cut_cells.is_empty() {
                continue;
            }
            let candidate = BalancedCut {
                cells: cut_cells,
                imbalance: left_count.abs_diff(right_count),
            };
            let key = (candidate.imbalance, candidate.cells.len());
            if direction_best
                .as_ref()
                .map_or(true, |old| key < (old.imbalance, old.cells.len()))
            {
                direction_best = Some(candidate);
            }
        }
        if let Some(candidate) = direction_best {
            let key = (candidate.cells.len(), candidate.imbalance);
            if best
                .as_ref()
                .map_or(true, |old| key < (old.cells.len(), old.imbalance))
            {
                best = Some(candidate);
            }
        }
    }
    best
}

// activeで指定されたマスだけの4近傍連結成分を取り出す。
// 抽出対象外の木を経由する実際の延焼は、別途行う焼失確率の計算に含まれる。
fn split_tree_sets(board: &Board, mut active: Vec<bool>) -> Vec<Vec<usize>> {
    let mut sets = Vec::new();
    for start in 0..active.len() {
        if !active[start] {
            continue;
        }
        active[start] = false;
        let mut cells = vec![start];
        let mut cursor = 0;
        while cursor < cells.len() {
            for &next in &board.neighbors[cells[cursor]] {
                if next != NONE && active[next] {
                    active[next] = false;
                    cells.push(next);
                }
            }
            cursor += 1;
        }
        sets.push(cells);
    }
    sets
}

// 乾燥度4以上の木を母集団とし、その上位high_risk_percent%を抽出する。
// 件数は切り捨て、同確率なら行・列の順。無視した木も実盤面には残す。
fn high_risk_mask(board: &Board, estimate: &BurnEstimate, high_risk_percent: usize) -> Vec<bool> {
    assert!(high_risk_percent <= 100);
    let mut trees: Vec<usize> = (0..board.n * board.n)
        .filter(|&v| board.dryness[v] > 3)
        .collect();
    trees.sort_unstable_by(|&a, &b| {
        estimate.burn_probabilities[b / board.n][b % board.n]
            .total_cmp(&estimate.burn_probabilities[a / board.n][a % board.n])
            .then_with(|| a.cmp(&b))
    });
    let keep = trees.len() * high_risk_percent / 100;
    let mut active = vec![false; board.n * board.n];
    for &v in &trees[..keep] {
        active[v] = true;
    }
    active
}

fn highest_expected_component(
    board: &Board,
    active: Vec<bool>,
    estimate: &BurnEstimate,
) -> Option<Vec<usize>> {
    let mut best: Option<(Vec<usize>, f64)> = None;
    for cells in split_tree_sets(board, active) {
        let expected: f64 = cells
            .iter()
            .map(|&v| estimate.burn_probabilities[v / board.n][v % board.n])
            .sum();
        if best.as_ref().map_or(true, |(_, old)| expected > *old) {
            best = Some((cells, expected));
        }
    }
    best.map(|(cells, _)| cells)
}

fn apply_cut(dryness: &mut [Vec<u32>], cut: &BalancedCut, cuts: &mut Vec<usize>, k: usize) -> bool {
    if cut.cells.len() > k - cuts.len() {
        return false;
    }
    let n = dryness.len() - 2;
    for &v in &cut.cells {
        assert!(dryness[v / n + 1][v % n + 1] > 0);
        dryness[v / n + 1][v % n + 1] = 0;
        cuts.push(v);
    }
    true
}

fn solve(dryness: &mut [Vec<u32>], k: usize, high_risk_percent: usize) -> Vec<usize> {
    let mut cuts = Vec::new();
    while cuts.len() < k {
        let estimate = estimate_expected_burned(dryness, EXPECTATION_SAMPLES);
        let board = Board::new(dryness);
        let active = high_risk_mask(&board, &estimate, high_risk_percent);
        let Some(cells) = highest_expected_component(&board, active, &estimate) else {
            break;
        };
        let Some(cut) = find_balanced_cut(&cells, board.n) else {
            break;
        };
        if !apply_cut(dryness, &cut, &mut cuts, k) {
            break;
        }
    }
    cuts
}

fn solve_best(dryness: &[Vec<u32>], k: usize) -> Vec<usize> {
    let mut best_expected = f64::INFINITY;
    let mut best_percent = usize::MAX;
    let mut best_cuts = Vec::new();
    for percent in HIGH_RISK_PERCENTS {
        // 候補ごとに初期盤面から開始し、最後の切断後の期待値で比較する。
        let mut board = dryness.to_vec();
        let cuts = solve(&mut board, k, percent);
        let expected = estimate_expected_burned(&board, EXPECTATION_SAMPLES).expected_burned;
        if expected
            .total_cmp(&best_expected)
            .then_with(|| percent.cmp(&best_percent))
            .is_lt()
        {
            best_expected = expected;
            best_percent = percent;
            best_cuts = cuts;
        }
    }
    best_cuts
}

fn main() {
    let mut input = String::new();
    io::stdin().read_to_string(&mut input).unwrap();
    let mut tokens = input.split_whitespace();
    let n: usize = tokens.next().unwrap().parse().unwrap();
    let k: usize = tokens.next().unwrap().parse().unwrap();
    // 外周を0で埋め、盤面外の乾燥度は合計に含めない。
    let mut dryness = vec![vec![0_u32; n + 2]; n + 2];
    for row in dryness.iter_mut().skip(1).take(n) {
        let digits = tokens.next().unwrap().as_bytes();
        for j in 0..n {
            row[j + 1] = (digits[j] - b'0') as u32;
        }
    }

    let cuts = solve_best(&dryness, k);
    println!("{}", cuts.len());
    for v in cuts {
        println!("{} {}", v / n, v % n);
    }
}