use std::collections::BinaryHeap;
use std::io::{self, BufRead, Write};
fn parse_first_line(line: &str) -> (i32, i32) {
let mut parts = line.split_whitespace();
let n = parts.next().unwrap().parse().unwrap();
let k = parts.next().unwrap().parse().unwrap();
(n, k)
}
fn parse_second_line(line: &str) -> Vec<i32> {
let parts = line.split_whitespace();
parts.map(|part| part.parse().unwrap()).collect()
}
// fn balance_nums(nums: &mut Vec<i32>, k: i32, acc: i32, index: usize) -> i32 {
// let mut new_acc = acc;
// if index > 0 {
// if nums[index] + k < nums[index - 1] {
// let to_add = nums[index - 1] - k - nums[index];
// nums[index] += to_add;
// new_acc = balance_nums(nums, k, new_acc + to_add, index - 1);
// } else if nums[index - 1] + k < nums[index] {
// let to_add = nums[index] - k - nums[index - 1];
// nums[index - 1] += to_add;
// new_acc = balance_nums(nums, k, new_acc + to_add, index - 1);
// }
// }
// if index < nums.len() - 1 {
// if nums[index] + k < nums[index + 1] {
// let to_add = nums[index + 1] - k - nums[index];
// nums[index] += to_add;
// new_acc = balance_nums(nums, k, new_acc + to_add, index + 1);
// } else if nums[index + 1] + k < nums[index] {
// let to_add = nums[index] - k - nums[index + 1];
// nums[index + 1] += to_add;
// new_acc = balance_nums(nums, k, new_acc + to_add, index + 1);
// }
// }
// new_acc
// }
// fn index_of_max_element(nums: &[i32]) -> usize {
// let mut max_index = 0;
// let mut max_value = nums[0];
// for (index, value) in nums.iter().enumerate() {
// if *value > max_value {
// max_index = index;
// max_value = *value;
// }
// }
// max_index
// }
fn main() {
let stdin = io::stdin().lock();
let mut stdout = io::stdout().lock();
let mut lines = stdin.lines();
let first_line = lines.next().unwrap().unwrap();
let second_line = lines.next().unwrap().unwrap();
let (_n, k) = parse_first_line(&first_line);
let mut nums = parse_second_line(&second_line);
let mut heap = BinaryHeap::new();
for (i, x) in nums.iter().enumerate() {
heap.push((*x, i));
}
let mut result = 0;
while !heap.is_empty() {
let (_height, index) = heap.pop().unwrap();
if index > 0 && nums[index - 1] + k < nums[index] {
let to_add = nums[index] - k - nums[index - 1];
nums[index - 1] += to_add;
result += to_add;
heap.push((nums[index - 1], index - 1));
}
if index < nums.len() - 1 && nums[index + 1] + k < nums[index] {
let to_add = nums[index] - k - nums[index + 1];
nums[index + 1] += to_add;
result += to_add;
heap.push((nums[index + 1], index + 1));
}
}
writeln!(stdout, "{}", result).expect("write stdout");
}
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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | use std::collections::BinaryHeap; use std::io::{self, BufRead, Write}; fn parse_first_line(line: &str) -> (i32, i32) { let mut parts = line.split_whitespace(); let n = parts.next().unwrap().parse().unwrap(); let k = parts.next().unwrap().parse().unwrap(); (n, k) } fn parse_second_line(line: &str) -> Vec<i32> { let parts = line.split_whitespace(); parts.map(|part| part.parse().unwrap()).collect() } // fn balance_nums(nums: &mut Vec<i32>, k: i32, acc: i32, index: usize) -> i32 { // let mut new_acc = acc; // if index > 0 { // if nums[index] + k < nums[index - 1] { // let to_add = nums[index - 1] - k - nums[index]; // nums[index] += to_add; // new_acc = balance_nums(nums, k, new_acc + to_add, index - 1); // } else if nums[index - 1] + k < nums[index] { // let to_add = nums[index] - k - nums[index - 1]; // nums[index - 1] += to_add; // new_acc = balance_nums(nums, k, new_acc + to_add, index - 1); // } // } // if index < nums.len() - 1 { // if nums[index] + k < nums[index + 1] { // let to_add = nums[index + 1] - k - nums[index]; // nums[index] += to_add; // new_acc = balance_nums(nums, k, new_acc + to_add, index + 1); // } else if nums[index + 1] + k < nums[index] { // let to_add = nums[index] - k - nums[index + 1]; // nums[index + 1] += to_add; // new_acc = balance_nums(nums, k, new_acc + to_add, index + 1); // } // } // new_acc // } // fn index_of_max_element(nums: &[i32]) -> usize { // let mut max_index = 0; // let mut max_value = nums[0]; // for (index, value) in nums.iter().enumerate() { // if *value > max_value { // max_index = index; // max_value = *value; // } // } // max_index // } fn main() { let stdin = io::stdin().lock(); let mut stdout = io::stdout().lock(); let mut lines = stdin.lines(); let first_line = lines.next().unwrap().unwrap(); let second_line = lines.next().unwrap().unwrap(); let (_n, k) = parse_first_line(&first_line); let mut nums = parse_second_line(&second_line); let mut heap = BinaryHeap::new(); for (i, x) in nums.iter().enumerate() { heap.push((*x, i)); } let mut result = 0; while !heap.is_empty() { let (_height, index) = heap.pop().unwrap(); if index > 0 && nums[index - 1] + k < nums[index] { let to_add = nums[index] - k - nums[index - 1]; nums[index - 1] += to_add; result += to_add; heap.push((nums[index - 1], index - 1)); } if index < nums.len() - 1 && nums[index + 1] + k < nums[index] { let to_add = nums[index] - k - nums[index + 1]; nums[index + 1] += to_add; result += to_add; heap.push((nums[index + 1], index + 1)); } } writeln!(stdout, "{}", result).expect("write stdout"); } |
English