#1195078 (Rust) No.5025 Forest Fire

• 問題: No.5025 Forest Fire
• ユーザー: shim0
• 提出日時: 2026-10-24 16:44:09
• 言語: Rust (1.99.0 + proconio + num + itertools + ACL)
• 結果: AC
• 実行時間: 1,964 ms / 2,000 ms
• コード長: 15,512 bytes
• コンパイル時間: 4,963 ms
• コンパイル使用メモリ: 208,888 KB
• 実行使用メモリ: 9,924 KB
• スコア: 164,189,171
• 最終ジャッジ日時: 2026-10-24 16:47:03
• 合計ジャッジ時間: 161,525 ms
• ジャッジサーバーID （参考情報）: judge6_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 TraceRow {
    elapsed: f64,
    attempt: u64,
    line: usize,
    axis: usize,
    from_axis: usize,
    move_kind: usize,
    from: usize,
    to: usize,
    candidate: usize,
    current: usize,
    best: usize,
    temperature: f64,
    accepted: bool,
}

fn write_trace(
    path: &std::path::Path,
    initial_elapsed: f64,
    initial: usize,
    rows: &[TraceRow],
) -> std::io::Result<()> {
    use std::io::Write;
    let mut out = std::io::BufWriter::new(std::fs::File::create(path)?);
    writeln!(
        out,
        "event,elapsed_s,attempt,line_id,axis,from,to,candidate_cost,current_cost,best_cost,temperature,accepted,move_kind,from_axis,move_accepted,move_rejected,accept_rate,reject_rate"
    )?;
    let t = (initial_elapsed / TIME_LIMIT).min(1.0);
    let temperature = 10000.0_f64.powf(1.0 - t) * 1e2_f64.powf(t);
    writeln!(
        out,
        "initial,{initial_elapsed:.9},0,,,,,{initial},{initial},{initial},{temperature:.9},,,,,,,"
    )?;
    let mut counts = [[0u64; 2]; 2];
    for r in rows {
        counts[r.move_kind][if r.accepted { 0 } else { 1 }] += 1;
        let [yes, no] = counts[r.move_kind];
        let total = (yes + no) as f64;
        writeln!(
            out,
            "candidate,{:.9},{},{},{},{},{},{},{},{},{:.9},{},{},{},{},{},{:.9},{:.9}",
            r.elapsed,
            r.attempt,
            r.line,
            if r.axis == 0 { "row" } else { "column" },
            r.from,
            r.to,
            r.candidate,
            r.current,
            r.best,
            r.temperature,
            r.accepted,
            if r.move_kind == 0 { "shift" } else { "rotate" },
            if r.from_axis == 0 { "row" } else { "column" },
            yes,
            no,
            yes as f64 / total,
            no as f64 / total
        )?;
    }
    out.flush()
}

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,
}

// Choose a balanced cut in the new orientation; ties retain the first line.
fn rotated_position(group: &[usize], n: usize, axis: usize) -> usize {
    let coordinate = |v: usize| if axis == 0 { v / n } else { v % n };
    let mut best = None;
    let mut best_count = usize::MAX;
    for line in 0..n {
        let count = group.iter().filter(|&&v| coordinate(v) == line).count();
        if count > 0 && count < best_count && is_balanced(group, n, axis, line) {
            best = Some(line);
            best_count = count;
        }
    }
    best.unwrap_or_else(|| {
        let low = group.iter().map(|&v| coordinate(v)).min().unwrap();
        let high = group.iter().map(|&v| coordinate(v)).max().unwrap();
        (low + high) / 2
    })
}

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();
    // Opt in locally: ANNEAL_LOG=/path/to/run.csv. No file I/O during search.
    let trace_path = std::env::var_os("ANNEAL_LOG").map(std::path::PathBuf::from);
    let mut trace = trace_path.as_ref().map(|_| Vec::with_capacity(100_000));
    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;
    let mut move_counts = [[0u64; 3]; 2];
    let initial_elapsed = started.elapsed().as_secs_f64();
    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 old_axis = lines[i].axis;
        let move_kind = if rng.index(10) < 8 { 0 } else { 1 };
        let new_axis = if move_kind == 0 {
            old_axis
        } else {
            1 - old_axis
        };
        let next = if move_kind == 0 {
            match rng.index(2) {
                0 if old > 0 => old - 1,
                1 if old + 1 < n => old + 1,
                _ => {
                    move_counts[move_kind][2] += 1;
                    continue;
                }
            }
        } else {
            rotated_position(&lines[i].group, n, new_axis)
        };
        lines[i].axis = new_axis;
        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 elapsed = started.elapsed().as_secs_f64();
        let t = (elapsed / TIME_LIMIT).min(1.0);
        let temperature = 10000.0_f64.powf(1.0 - t) * 1e2_f64.powf(t);
        let did_accept = delta >= 0.0 || rng.unit() < (delta / temperature).exp();
        move_counts[move_kind][if did_accept { 0 } else { 1 }] += 1;
        if did_accept {
            accepted += 1;
            current = candidate;
            if current < best_cost {
                best_cost = current;
                best.clone_from(&output);
            }
        } else {
            lines[i].position = old;
            lines[i].axis = old_axis;
        }
        if let Some(trace) = &mut trace {
            trace.push(TraceRow {
                elapsed,
                attempt: attempts,
                line: i,
                axis: new_axis,
                from_axis: old_axis,
                move_kind,
                from: old,
                to: next,
                candidate,
                current,
                best: best_cost,
                temperature,
                accepted: did_accept,
            });
        }
    }
    eprintln!(
        "initial_cost: {initial_cost}, best_cost: {best_cost}, attempts: {attempts}, accepted: {accepted}"
    );
    for (kind, counts) in ["shift", "rotate"].iter().zip(move_counts) {
        let [yes, no, skipped] = counts;
        let total = (yes + no) as f64;
        let accept_rate = if total > 0.0 { yes as f64 / total } else { 0.0 };
        let reject_rate = if total > 0.0 { no as f64 / total } else { 0.0 };
        eprintln!(
            "{kind}: accepted={yes}, rejected={no}, skipped={skipped}, accept_rate={accept_rate:.6}, reject_rate={reject_rate:.6}"
        );
    }
    eprintln!("remaining_cut_budget: {}", k - best.len());
    println!("{}", best.len());
    for v in best {
        println!("{} {}", v / n, v % n);
    }
    if let Some(path) = trace_path {
        match write_trace(
            &path,
            initial_elapsed,
            initial_cost,
            trace.as_deref().unwrap(),
        ) {
            Ok(()) => eprintln!("annealing_log: {}", path.display()),
            Err(error) => eprintln!("failed to write annealing log {}: {error}", path.display()),
        }
    }
}