From e04d041d2680fd0dc7cb752e4c42a817bf894c1d Mon Sep 17 00:00:00 2001 From: EFanZh Date: Tue, 14 May 2024 10:04:36 +0800 Subject: [PATCH] Add problem 1825: Finding MK Average --- src/lib.rs | 1 + .../btree_map.rs | 235 ++++++++++++++++++ src/problem_1825_finding_mk_average/mod.rs | 131 ++++++++++ 3 files changed, 367 insertions(+) create mode 100644 src/problem_1825_finding_mk_average/btree_map.rs create mode 100644 src/problem_1825_finding_mk_average/mod.rs diff --git a/src/lib.rs b/src/lib.rs index 405f9067..48fc68cc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1363,6 +1363,7 @@ pub mod problem_1818_minimum_absolute_sum_difference; pub mod problem_1819_number_of_different_subsequences_gcds; pub mod problem_1822_sign_of_the_product_of_an_array; pub mod problem_1824_minimum_sideway_jumps; +pub mod problem_1825_finding_mk_average; pub mod problem_1827_minimum_operations_to_make_the_array_increasing; pub mod problem_1829_maximum_xor_for_each_query; pub mod problem_1832_check_if_the_sentence_is_pangram; diff --git a/src/problem_1825_finding_mk_average/btree_map.rs b/src/problem_1825_finding_mk_average/btree_map.rs new file mode 100644 index 00000000..5c620eab --- /dev/null +++ b/src/problem_1825_finding_mk_average/btree_map.rs @@ -0,0 +1,235 @@ +// ------------------------------------------------------ snip ------------------------------------------------------ // + +use std::collections::btree_map::{Entry, OccupiedEntry}; +use std::collections::{BTreeMap, VecDeque}; +use std::mem; +use std::num::NonZeroU64; + +#[derive(Default)] +struct MultiSet(BTreeMap); + +impl MultiSet { + fn insert(&mut self, num: u32) { + match self.0.entry(num) { + Entry::Vacant(entry) => { + entry.insert(0); + } + Entry::Occupied(entry) => *entry.into_mut() += 1, + } + } + + fn remove_entry(entry: OccupiedEntry) { + if *entry.get() == 0 { + entry.remove(); + } else { + *entry.into_mut() -= 1; + } + } + + fn remove(&mut self, num: u32) -> bool { + match self.0.entry(num) { + Entry::Vacant(_) => false, + Entry::Occupied(entry) => { + Self::remove_entry(entry); + + true + } + } + } + + fn replace_first(&mut self, num: u32) -> Option { + let entry = self.0.first_entry().unwrap(); + let key = *entry.key(); + + #[allow(clippy::if_then_some_else_none)] // False positive. + if num > key { + Self::remove_entry(entry); + self.insert(num); + + Some(key) + } else { + None + } + } + + fn replace_last(&mut self, num: u32) -> Option { + let entry = self.0.last_entry().unwrap(); + let key = *entry.key(); + + #[allow(clippy::if_then_some_else_none)] // False positive. + if num < key { + Self::remove_entry(entry); + self.insert(num); + + Some(key) + } else { + None + } + } +} + +enum Inner { + Start { + nums: Vec, + required: u32, + k: u32, + }, + Running { + queue: VecDeque, + low: MultiSet, + middle: MultiSet, + high: MultiSet, + middle_sum: u64, + middle_length: NonZeroU64, + }, +} + +pub struct MKAverage { + inner: Inner, +} + +impl MKAverage { + fn new(m: i32, k: i32) -> Self { + let m = m as u32; + let k = k as u32; + + Self { + inner: Inner::Start { + nums: Vec::with_capacity(m as _), + required: m, + k, + }, + } + } + + fn count_nums(nums: &[u32]) -> MultiSet { + let mut result = MultiSet::default(); + + for &num in nums { + result.insert(num); + } + + result + } + + fn add_element(&mut self, num: i32) { + let num = num as u32; + + match &mut self.inner { + Inner::Start { nums, required, k } => { + nums.push(num); + + *required -= 1; + + if *required == 0 { + let m = nums.len(); + let k = *k as usize; + let middle_length = m - k - k; + let mut buffer = nums.clone().into_boxed_slice(); + + buffer.select_nth_unstable(k - 1).2.select_nth_unstable(middle_length); + + let (low, rest) = buffer.split_at(k); + let (middle, high) = rest.split_at(middle_length); + + let middle_sum = middle.iter().fold(0, |sum, &value| sum + u64::from(value)); + let queue = VecDeque::from(mem::take(nums)); + let low = Self::count_nums(low); + let middle = Self::count_nums(middle); + let high = Self::count_nums(high); + + drop(mem::replace( + &mut self.inner, + Inner::Running { + queue, + low, + middle, + high, + middle_sum, + middle_length: NonZeroU64::new(middle_length as _).unwrap(), + }, + )); + } + } + Inner::Running { + queue, + low, + middle, + high, + middle_sum, + .. + } => { + let old_num = queue.pop_front().unwrap(); + + queue.push_back(num); + + if old_num != num { + if low.remove(old_num) { + let new_middle = high.replace_first(num).unwrap_or(num); + + low.insert(middle.replace_first(new_middle).map_or(new_middle, |middle_min| { + *middle_sum -= u64::from(middle_min); + *middle_sum += u64::from(new_middle); + + middle_min + })); + } else if high.remove(old_num) { + let new_middle = low.replace_last(num).unwrap_or(num); + + high.insert(middle.replace_last(new_middle).map_or(new_middle, |middle_max| { + *middle_sum -= u64::from(middle_max); + *middle_sum += u64::from(new_middle); + + middle_max + })); + } else { + middle.remove(old_num); + *middle_sum -= u64::from(old_num); + + let new_middle = low + .replace_last(num) + .unwrap_or_else(|| high.replace_first(num).unwrap_or(num)); + + middle.insert(new_middle); + *middle_sum += u64::from(new_middle); + } + } + } + } + } + + fn calculate_mk_average(&self) -> i32 { + match self.inner { + Inner::Start { .. } => -1, + Inner::Running { + middle_sum, + middle_length, + .. + } => (middle_sum / middle_length) as _, + } + } +} + +// ------------------------------------------------------ snip ------------------------------------------------------ // + +impl super::MKAverage for MKAverage { + fn new(m: i32, k: i32) -> Self { + Self::new(m, k) + } + + fn add_element(&mut self, num: i32) { + self.add_element(num); + } + + fn calculate_mk_average(&self) -> i32 { + self.calculate_mk_average() + } +} + +#[cfg(test)] +mod tests { + #[test] + fn test_solution() { + super::super::tests::run::(); + } +} diff --git a/src/problem_1825_finding_mk_average/mod.rs b/src/problem_1825_finding_mk_average/mod.rs new file mode 100644 index 00000000..64f2053b --- /dev/null +++ b/src/problem_1825_finding_mk_average/mod.rs @@ -0,0 +1,131 @@ +pub mod btree_map; + +pub trait MKAverage { + fn new(m: i32, k: i32) -> Self; + fn add_element(&mut self, num: i32); + fn calculate_mk_average(&self) -> i32; +} + +#[cfg(test)] +mod tests { + use super::MKAverage; + + enum Operation { + AddElement(i32), + CalculateMkAverage(i32), + } + + type TestCase<'a> = (i32, i32, &'a [Operation]); + + const EXTRA_TEST_CASE_1: TestCase = ( + 6, + 2, + &[ + Operation::AddElement(1), + Operation::AddElement(1), + Operation::AddElement(2), + Operation::AddElement(2), + Operation::AddElement(2), + Operation::AddElement(2), + Operation::AddElement(2), + Operation::AddElement(2), + Operation::AddElement(2), + Operation::CalculateMkAverage(2), + ], + ); + + const EXTRA_TEST_CASE_2: TestCase = ( + 3, + 1, + &[ + Operation::AddElement(8), + Operation::AddElement(7), + Operation::AddElement(7), + Operation::AddElement(2), + Operation::AddElement(3), + Operation::AddElement(1), + Operation::CalculateMkAverage(2), + ], + ); + + pub fn run() { + let test_cases = [ + ( + 3, + 1, + &[ + Operation::AddElement(3), + Operation::AddElement(1), + Operation::CalculateMkAverage(-1), + Operation::AddElement(10), + Operation::CalculateMkAverage(3), + Operation::AddElement(5), + Operation::AddElement(5), + Operation::AddElement(5), + Operation::CalculateMkAverage(5), + ] as &[_], + ), + ( + 3, + 1, + &[ + Operation::AddElement(58916), + Operation::AddElement(61899), + Operation::CalculateMkAverage(-1), + Operation::AddElement(85406), + Operation::AddElement(49757), + Operation::CalculateMkAverage(61899), + Operation::AddElement(27520), + Operation::AddElement(12303), + Operation::CalculateMkAverage(27520), + Operation::AddElement(63945), + ], + ), + ( + 3, + 1, + &[ + Operation::AddElement(3716), + Operation::AddElement(51094), + Operation::CalculateMkAverage(-1), + Operation::AddElement(56724), + Operation::AddElement(79619), + Operation::CalculateMkAverage(56724), + Operation::AddElement(99914), + Operation::AddElement(277), + Operation::CalculateMkAverage(79619), + Operation::AddElement(91205), + ], + ), + ( + 3, + 1, + &[ + Operation::AddElement(17612), + Operation::AddElement(74607), + Operation::CalculateMkAverage(-1), + Operation::AddElement(8272), + Operation::AddElement(33433), + Operation::CalculateMkAverage(33433), + Operation::AddElement(15456), + Operation::AddElement(64938), + Operation::CalculateMkAverage(33433), + Operation::AddElement(99741), + ], + ), + EXTRA_TEST_CASE_1, + EXTRA_TEST_CASE_2, + ]; + + for (m, k, operations) in test_cases { + let mut mk_average = M::new(m, k); + + for operation in operations { + match *operation { + Operation::AddElement(num) => mk_average.add_element(num), + Operation::CalculateMkAverage(expected) => assert_eq!(mk_average.calculate_mk_average(), expected), + } + } + } + } +}