use std::{
collections::BinaryHeap,
io::{self, Read},
};
#[derive(Debug, Eq, PartialEq)]
struct Field {
index: usize,
height: usize,
}
impl Ord for Field {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.height
.cmp(&other.height)
.then_with(|| self.index.cmp(&other.index))
}
}
impl PartialOrd for Field {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
fn main() {
let mut input = String::new();
io::stdin()
.read_to_string(&mut input)
.expect("Failed to read input");
let mut scan = input.split_whitespace();
let mut next_num = || scan.next().unwrap().parse::<usize>().unwrap();
let n = next_num();
let k = next_num();
let mut heap: BinaryHeap<Field> = BinaryHeap::with_capacity(n);
let mut heights: Vec<usize> = Vec::with_capacity(n);
for i in 0..n {
let height = next_num();
heap.push(Field { index: i, height });
heights.push(height);
}
let mut result: usize = 0;
while let Some(field) = heap.pop() {
let Field {
index,
height: element_height,
} = field;
// skip outdated element
if element_height != heights[index] {
continue;
}
let left = index.checked_sub(1);
let right = (index + 1 < n).then_some(index + 1);
for neighbor_index in [left, right].into_iter().flatten() {
if let Some(neighbor_height) = heights.get_mut(neighbor_index) {
if *neighbor_height + k < element_height {
let diff = element_height - k - *neighbor_height;
*neighbor_height += diff;
result += diff;
heap.push(Field {
index: neighbor_index,
height: *neighbor_height,
})
}
}
}
}
println!("{}", result);
}
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | use std::{ collections::BinaryHeap, io::{self, Read}, }; #[derive(Debug, Eq, PartialEq)] struct Field { index: usize, height: usize, } impl Ord for Field { fn cmp(&self, other: &Self) -> std::cmp::Ordering { self.height .cmp(&other.height) .then_with(|| self.index.cmp(&other.index)) } } impl PartialOrd for Field { fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { Some(self.cmp(other)) } } fn main() { let mut input = String::new(); io::stdin() .read_to_string(&mut input) .expect("Failed to read input"); let mut scan = input.split_whitespace(); let mut next_num = || scan.next().unwrap().parse::<usize>().unwrap(); let n = next_num(); let k = next_num(); let mut heap: BinaryHeap<Field> = BinaryHeap::with_capacity(n); let mut heights: Vec<usize> = Vec::with_capacity(n); for i in 0..n { let height = next_num(); heap.push(Field { index: i, height }); heights.push(height); } let mut result: usize = 0; while let Some(field) = heap.pop() { let Field { index, height: element_height, } = field; // skip outdated element if element_height != heights[index] { continue; } let left = index.checked_sub(1); let right = (index + 1 < n).then_some(index + 1); for neighbor_index in [left, right].into_iter().flatten() { if let Some(neighbor_height) = heights.get_mut(neighbor_index) { if *neighbor_height + k < element_height { let diff = element_height - k - *neighbor_height; *neighbor_height += diff; result += diff; heap.push(Field { index: neighbor_index, height: *neighbor_height, }) } } } } println!("{}", result); } |
English