#1195048 (Rust) No.5025 Forest Fire

• 問題: No.5025 Forest Fire
• ユーザー: shim0
• 提出日時: 2026-10-24 16:21:21
• 言語: Rust (1.99.0 + proconio + num + itertools + ACL)
• 結果: AC
• 実行時間: 1,986 ms / 2,000 ms
• コード長: 10,584 bytes
• コンパイル時間: 4,901 ms
• コンパイル使用メモリ: 204,968 KB
• 実行使用メモリ: 9,924 KB
• スコア: 161,563,808
• 最終ジャッジ日時: 2026-10-24 16:24:03
• 合計ジャッジ時間: 161,857 ms
• ジャッジサーバーID （参考情報）: judge4_0 / <nil>

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

ソースコード

use ac_library::Dsu;
use proconio::{input, marker::Bytes};
use std::time::Instant;

const TIME_LIMIT: f64 = 2.0;

struct Rng(u64);

impl Rng {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E3779B97F4A7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
        z ^ (z >> 31)
    }

    fn index(&mut self, len: usize) -> usize {
        let len = len as u64;
        let threshold = len.wrapping_neg() % len;
        loop {
            let x = self.next();
            if x >= threshold {
                return (x % len) as usize;
            }
        }
    }

    fn unit(&mut self) -> f64 {
        (self.next() >> 11) as f64 / ((1u64 << 53) as f64)
    }
}

struct Evaluator {
    n: usize,
    dryness: Vec<u8>,
    neighbors: Vec<Vec<usize>>,
    starts: Vec<usize>,
    // Each list is sorted by Manhattan distance, then row-major order.
    replacements: Vec<Vec<usize>>,
    component: Vec<usize>,
    sizes: Vec<usize>,
    queue: Vec<usize>,
}

impl Evaluator {
    fn new(n: usize, dryness: Vec<u8>) -> Self {
        let mut neighbors = vec![Vec::new(); n * n];
        for (v, adjacent) in neighbors.iter_mut().enumerate() {
            let (r, c) = (v / n, v % n);
            if r > 0 {
                adjacent.push(v - n);
            }
            if r + 1 < n {
                adjacent.push(v + n);
            }
            if c > 0 {
                adjacent.push(v - 1);
            }
            if c + 1 < n {
                adjacent.push(v + 1);
            }
        }
        let mut starts = Vec::new();
        let mut replacements = Vec::new();
        for r in (5..n).step_by(5) {
            for c in (5..n).step_by(5) {
                starts.push(r * n + c);
                let mut trees: Vec<_> = (0..n * n).filter(|&v| dryness[v] > 0).collect();
                trees.sort_unstable_by_key(|&v| (r.abs_diff(v / n) + c.abs_diff(v % n), v));
                replacements.push(trees);
            }
        }
        Self {
            n,
            dryness,
            neighbors,
            starts,
            replacements,
            component: vec![usize::MAX; n * n],
            sizes: Vec::with_capacity(n * n),
            queue: Vec::with_capacity(n * n),
        }
    }

    fn cost(&mut self, cut: &[bool]) -> usize {
        // A component's size equals the BFS area for every start inside it.
        // Compute it once, but count it separately for every start as specified.
        let mut total = 0;
        for threshold in [5, 7] {
            self.component.fill(usize::MAX);
            self.sizes.clear();
            for start in 0..self.n * self.n {
                if cut[start]
                    || self.dryness[start] < threshold
                    || self.component[start] != usize::MAX
                {
                    continue;
                }
                let id = self.sizes.len();
                self.queue.clear();
                self.queue.push(start);
                self.component[start] = id;
                let mut head = 0;
                while head < self.queue.len() {
                    let v = self.queue[head];
                    head += 1;
                    for &u in &self.neighbors[v] {
                        if !cut[u]
                            && self.dryness[u] >= threshold
                            && self.component[u] == usize::MAX
                        {
                            self.component[u] = id;
                            self.queue.push(u);
                        }
                    }
                }
                self.sizes.push(self.queue.len());
            }
            for (i, &original) in self.starts.iter().enumerate() {
                let start = if cut[original] {
                    self.replacements[i].iter().copied().find(|&v| !cut[v])
                } else {
                    Some(original)
                };
                if let Some(v) = start {
                    let id = self.component[v];
                    if id != usize::MAX {
                        total += self.sizes[id];
                    }
                }
            }
        }
        total
    }
}

const DRY_THRESHOLD: u8 = 5;

fn neighbors(v: usize, n: usize) -> impl Iterator<Item = usize> {
    let (r, c) = (v / n, v % n);
    [
        (r > 0).then(|| v - n),
        (r + 1 < n).then(|| v + n),
        (c > 0).then(|| v - 1),
        (c + 1 < n).then(|| v + 1),
    ]
    .into_iter()
    .flatten()
}

fn largest_group(active: &[bool], n: usize) -> Vec<usize> {
    let mut dsu = Dsu::new(n * n);
    for v in 0..n * n {
        if active[v] {
            for u in neighbors(v, n) {
                if active[u] {
                    dsu.merge(v, u);
                }
            }
        }
    }
    let mut best = Vec::new();
    for group in dsu.groups() {
        if active[group[0]] && group.len() > best.len() {
            best = group;
        }
    }
    best
}

// Test actual connected components after removing this row or column.
fn is_balanced(group: &[usize], n: usize, axis: usize, line: usize) -> bool {
    let coordinate = |v: usize| if axis == 0 { v / n } else { v % n };
    let mut unseen = vec![false; n * n];
    for &v in group {
        unseen[v] = coordinate(v) != line;
    }
    let mut stack = Vec::with_capacity(group.len());
    for &start in group {
        if !unseen[start] {
            continue;
        }
        unseen[start] = false;
        stack.push(start);
        let mut size = 0;
        while let Some(v) = stack.pop() {
            size += 1;
            if size * 2 > group.len() {
                return false;
            }
            for u in neighbors(v, n) {
                if unseen[u] {
                    unseen[u] = false;
                    stack.push(u);
                }
            }
        }
    }
    true
}

fn choose_cut(group: &[usize], n: usize) -> Option<Vec<usize>> {
    let mut best: Option<Vec<usize>> = None;
    // First encountered wins ties: rows, then columns, in ascending order.
    for axis in 0..2 {
        for line in 0..n {
            let cut: Vec<_> = group
                .iter()
                .copied()
                .filter(|&v| (if axis == 0 { v / n } else { v % n }) == line)
                .collect();
            if cut.is_empty() || best.as_ref().is_some_and(|b| cut.len() >= b.len()) {
                continue;
            }
            if is_balanced(group, n, axis, line) {
                best = Some(cut);
            }
        }
    }
    best
}

struct CutLine {
    group: Vec<usize>,
    axis: usize,
    position: usize,
}

fn initial_lines(dryness: &[u8], n: usize, k: usize) -> Vec<CutLine> {
    let mut active: Vec<_> = dryness.iter().map(|&d| d >= DRY_THRESHOLD).collect();
    let mut lines = Vec::new();
    let mut count = 0;
    while count < k {
        let group = largest_group(&active, n);
        if group.is_empty() {
            break;
        }
        let mut cut = choose_cut(&group, n)
            .expect("No straight cut satisfies the specified half-size condition");
        cut.sort_unstable();
        // Recover the first matching row/column in the same order as choose_cut.
        let (axis, position) = (0..2)
            .flat_map(|axis| (0..n).map(move |line| (axis, line)))
            .find(|&(axis, line)| {
                let cells: Vec<_> = group
                    .iter()
                    .copied()
                    .filter(|&v| (if axis == 0 { v / n } else { v % n }) == line)
                    .collect();
                cells == cut
            })
            .unwrap();
        for &v in cut.iter().take(k - count) {
            active[v] = false;
            count += 1;
        }
        lines.push(CutLine {
            group,
            axis,
            position,
        });
    }
    lines
}

fn materialize(lines: &[CutLine], n: usize, k: usize, cut: &mut [bool], output: &mut Vec<usize>) {
    cut.fill(false);
    output.clear();
    for line in lines {
        // Stored groups are row-major, so truncation cuts from the edge.
        for &v in &line.group {
            if output.len() == k {
                return;
            }
            let coordinate = if line.axis == 0 { v / n } else { v % n };
            if coordinate == line.position && !cut[v] {
                cut[v] = true;
                output.push(v);
            }
        }
    }
}

fn main() {
    let started = Instant::now();
    input! { n: usize, k: usize, rows: [Bytes; n] }
    let dryness: Vec<_> = rows
        .into_iter()
        .flat_map(|row| {
            assert_eq!(row.len(), n);
            row.into_iter().map(|d| d - b'0')
        })
        .collect();
    let mut lines = initial_lines(&dryness, n, k);
    let mut cut = vec![false; n * n];
    let mut output = Vec::with_capacity(k);
    materialize(&lines, n, k, &mut cut, &mut output);
    let mut evaluator = Evaluator::new(n, dryness);
    let mut current = evaluator.cost(&cut);
    let initial_cost = current;
    let mut best_cost = current;
    let mut best = output.clone();
    let mut rng = Rng(20261024);
    let mut attempts = 0u64;
    let mut accepted = 0u64;
    while !lines.is_empty() && started.elapsed().as_secs_f64() < TIME_LIMIT {
        attempts += 1;
        let i = rng.index(lines.len());
        let old = lines[i].position;
        let next = match rng.index(2) {
            0 if old > 0 => old - 1,
            1 if old + 1 < n => old + 1,
            _ => continue,
        };
        lines[i].position = next;
        materialize(&lines, n, k, &mut cut, &mut output);
        let candidate = evaluator.cost(&cut);
        let delta = current as f64 - candidate as f64;
        let t = (started.elapsed().as_secs_f64() / TIME_LIMIT).min(1.0);
        let temperature = 1e4_f64.powf(1.0 - t) * 1e2_f64.powf(t);
        if delta >= 0.0 || rng.unit() < (delta / temperature).exp() {
            accepted += 1;
            current = candidate;
            if current < best_cost {
                best_cost = current;
                best.clone_from(&output);
            }
        } else {
            lines[i].position = old;
        }
    }
    eprintln!(
        "initial_cost: {initial_cost}, best_cost: {best_cost}, attempts: {attempts}, accepted: {accepted}"
    );
    eprintln!("remaining_cut_budget: {}", k - best.len());
    println!("{}", best.len());
    for v in best {
        println!("{} {}", v / n, v % n);
    }
}