AtCoder Beginner Contest 384 E

https://atcoder.jp/contests/abc384/tasks/abc384_e

周りのマスをQueueに入れて、値が小さい順に取り出します。

// Takahashi is Slime 2
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// Matirx ////////////////////

type Point = (usize, usize);
type Matrix = Vec<Vec<u64>>;

fn read_matrix(H: usize) -> Matrix {
    (0..H).map(|_| read_vec::<u64>()).collect::<Vec<_>>()
}


//////////////////// BinaryHeap ////////////////////

use std::collections::BinaryHeap;

#[derive(Debug, Eq, PartialEq)]
struct Item {
    value: u64,
    pt: Point
}

impl Ord for Item {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        other.value.cmp(&self.value)
    }
}

impl PartialOrd for Item {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}


//////////////////// process ////////////////////

fn read_input() -> (u64, usize, usize, Matrix) {
    let v: Vec<usize> = read_vec();
    let (H, X) = (v[0], v[2] as u64);
    let w: Vec<usize> = read_vec();
    let (P, Q) = (w[0] - 1, w[1] - 1);
    let S = read_matrix(H);
    (X, P, Q, S)
}

fn neighbors(pt: Point, S: &Matrix) -> Vec<Point> {
    let H = S.len();
    let W = S[0].len();
    let (x, y) = pt;
    let mut neighs: Vec<Point> = vec![];
    if x > 0 {
        neighs.push((x - 1, y))
    }
    if x < H - 1 {
        neighs.push((x + 1, y))
    }
    if y > 0 {
        neighs.push((x, y - 1))
    }
    if y < W - 1 {
        neighs.push((x, y + 1))
    }
    neighs
}

fn F(X: u64, P: usize, Q: usize, S: Matrix) -> u64 {
    let H = S.len();
    let W = S[0].len();
    let mut queued: Vec<Vec<bool>> = vec![vec![false; W]; H];
    queued[P][Q] = true;
    let mut s = S[P][Q];
    let mut heap = BinaryHeap::new();
    for pt in neighbors((P, Q), &S).into_iter() {
        heap.push(Item { value: S[pt.0][pt.1], pt });
        queued[pt.0][pt.1] = true
    }
    
    while let Some(item) = heap.pop() {
        if item.value >= (s + X - 1) / X {
            break
        }
        s += item.value;
        for pt in neighbors(item.pt, &S) {
            if !queued[pt.0][pt.1] {
                heap.push(Item { value: S[pt.0][pt.1], pt });
                queued[pt.0][pt.1] = true
            }
        }
    }
    s
}

fn main() {
    let (X, P, Q, S) = read_input();
    println!("{}", F(X, P, Q, S))
}