From 41187ad27f0f6920ab9feb0fb32ad48ce81274a4 Mon Sep 17 00:00:00 2001 From: EFanZh Date: Sat, 14 Dec 2024 19:48:43 +0800 Subject: [PATCH] Add problem 2333: Minimum Sum of Squared Difference --- src/lib.rs | 1 + .../mod.rs | 47 ++++++ .../quick_select.rs | 158 ++++++++++++++++++ 3 files changed, 206 insertions(+) create mode 100644 src/problem_2333_minimum_sum_of_squared_difference/mod.rs create mode 100644 src/problem_2333_minimum_sum_of_squared_difference/quick_select.rs diff --git a/src/lib.rs b/src/lib.rs index 4761fd83..d3921b71 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1736,6 +1736,7 @@ pub mod problem_2326_spiral_matrix_iv; pub mod problem_2327_number_of_people_aware_of_a_secret; pub mod problem_2328_number_of_increasing_paths_in_a_grid; pub mod problem_2331_evaluate_boolean_binary_tree; +pub mod problem_2333_minimum_sum_of_squared_difference; pub mod problem_2335_minimum_amount_of_time_to_fill_cups; pub mod problem_2336_smallest_number_in_infinite_set; pub mod problem_2337_move_pieces_to_obtain_a_string; diff --git a/src/problem_2333_minimum_sum_of_squared_difference/mod.rs b/src/problem_2333_minimum_sum_of_squared_difference/mod.rs new file mode 100644 index 00000000..5b566101 --- /dev/null +++ b/src/problem_2333_minimum_sum_of_squared_difference/mod.rs @@ -0,0 +1,47 @@ +pub mod quick_select; + +pub trait Solution { + fn min_sum_square_diff(nums1: Vec, nums2: Vec, k1: i32, k2: i32) -> i64; +} + +#[cfg(test)] +mod tests { + use super::Solution; + + pub fn run() { + let test_cases = [ + ((&[1, 2, 3, 4] as &[_], &[2, 10, 20, 19] as &[_], 0, 0), 579), + ((&[1, 4, 10, 12], &[5, 8, 6, 9], 1, 1), 43), + ((&[18, 4, 8, 19, 13, 8], &[18, 11, 8, 2, 13, 15], 16, 8), 17), + ( + (&[7, 11, 4, 19, 11, 5, 6, 1, 8], &[4, 7, 6, 16, 12, 9, 10, 2, 10], 3, 6), + 27, + ), + ( + ( + &[ + 32, 17, 74, 75, 84, 56, 52, 29, 67, 71, 6, 1, 92, 51, 9, 100, 49, 95, 24, 42, 34, 36, 19, 53, + 65, 13, 98, 86, 70, 29, 89, 32, 46, 89, 17, 21, 56, 68, 98, 75, 44, 31, 56, 6, 74, 23, 48, 10, + 40, 55, 4, 99, 35, 42, 15, 42, 97, 10, 72, 75, 14, 49, 1, 61, 6, 37, 41, 31, 23, 66, 15, 91, + 30, 96, 44, 23, 31, 28, 24, 75, 68, 52, 7, 21, 7, 54, 88, 7, 52, 98, 25, 15, 58, 31, 33, 0, 55, + 89, 45, 5, + ], + &[ + 21, 91, 81, 85, 20, 41, 90, 71, 64, 31, 56, 22, 16, 22, 15, 30, 13, 78, 30, 12, 44, 0, 56, 32, + 10, 44, 47, 84, 57, 43, 30, 86, 75, 77, 47, 86, 57, 45, 33, 62, 73, 17, 92, 12, 7, 88, 63, 6, + 67, 91, 37, 49, 60, 57, 89, 89, 21, 15, 4, 83, 7, 1, 89, 30, 25, 27, 12, 63, 60, 31, 96, 33, + 28, 97, 53, 75, 19, 100, 78, 52, 96, 91, 92, 91, 90, 38, 7, 88, 19, 86, 27, 38, 37, 27, 61, 71, + 52, 23, 4, 41, + ], + 138, + 3351, + ), + 0, + ), + ]; + + for ((nums1, nums2, k1, k2), expected) in test_cases { + assert_eq!(S::min_sum_square_diff(nums1.to_vec(), nums2.to_vec(), k1, k2), expected); + } + } +} diff --git a/src/problem_2333_minimum_sum_of_squared_difference/quick_select.rs b/src/problem_2333_minimum_sum_of_squared_difference/quick_select.rs new file mode 100644 index 00000000..7d82a235 --- /dev/null +++ b/src/problem_2333_minimum_sum_of_squared_difference/quick_select.rs @@ -0,0 +1,158 @@ +pub struct Solution; + +// ------------------------------------------------------ snip ------------------------------------------------------ // + +struct Random(usize); + +impl Random { + fn next(&mut self, seed: usize) -> usize { + self.0 = self.0.wrapping_add(seed).wrapping_mul(N); + + self.0 + } +} + +impl Solution { + fn three_way_partition(nums: &mut [u32], key: u32) -> (usize, usize) { + let n = nums.len(); + let mut equal_start = 0; + let mut unchecked_start = 0; + let mut greater_start = n; + + 'outer: while unchecked_start < greater_start { + let left_num = nums[unchecked_start]; + + if left_num <= key { + if left_num < key { + if equal_start < unchecked_start { + nums[equal_start] = left_num; + nums[unchecked_start] = key; + } + + equal_start += 1; + } + } else { + loop { + greater_start -= 1; + + if unchecked_start < greater_start { + let right_num = nums[greater_start]; + + if right_num <= key { + nums[greater_start] = left_num; + + if right_num < key { + nums[equal_start] = right_num; + + if equal_start < unchecked_start { + nums[unchecked_start] = key; + } + + equal_start += 1; + } else { + nums[unchecked_start] = key; + } + + break; + } + } else { + break 'outer; + } + } + } + + unchecked_start += 1; + } + + (equal_start, unchecked_start) + } + + pub fn min_sum_square_diff(nums1: Vec, nums2: Vec, k1: i32, k2: i32) -> i64 { + let mut nums1 = nums1; + + nums1 + .iter_mut() + .zip(nums2) + .for_each(|(target, source)| *target = target.abs_diff(source) as _); + + let mut nums = nums1.into_iter().map(|x| x as u32).collect::>(); + let nums = &mut *nums; + let n = nums.len(); + let k = u64::from((k1 + k2) as u32); + let mut random = Random::<0x_9E37_79B9>(0); + let mut left = 0; + let mut right = n; + let mut right_sum = 0; + + while left < right { + let window = &mut nums[left..right]; + let window_len = window.len(); + let middle_num = window[random.next(window_len) % window_len]; + let (mut middle_start, mut middle_end) = Self::three_way_partition(window, middle_num); + + middle_start += left; + middle_end += left; + + let (upper_bound, new_left, new_right) = if middle_end - middle_start < 2 { + ( + if middle_start == left { + nums.get(middle_start.wrapping_sub(1)).copied().unwrap_or(0) + } else { + nums[left..middle_start].iter().copied().max().unwrap() + }, + middle_start + 1, + middle_start, + ) + } else { + (middle_num, middle_end, middle_start + 1) + }; + + let candidate_right_sum = nums[middle_end..right].iter().fold( + u64::from(middle_num) * (middle_end - new_right) as u64 + right_sum, + |sum, &x| sum + u64::from(x), + ); + + let extra = candidate_right_sum - u64::from(upper_bound) * (n - new_right) as u64; + + if k < extra { + left = new_left; + } else { + right_sum = candidate_right_sum; + right = new_right; + } + } + + nums.get(left.wrapping_sub(1)).map_or(0, |&upper_bound| { + let upper_bound = u64::from(upper_bound); + let mut right_len = (n - left) as u64; + let remaining = k - (right_sum - upper_bound * right_len); + + right_len += 1; + + let quotient = remaining / right_len; + let remainder = remaining % right_len; + + nums[..left - 1].iter().fold( + (upper_bound - quotient - 1).pow(2) * remainder + + (upper_bound - quotient).pow(2) * (right_len - remainder), + |sum, &x| sum + u64::from(x).pow(2), + ) + }) as _ + } +} + +// ------------------------------------------------------ snip ------------------------------------------------------ // + +impl super::Solution for Solution { + fn min_sum_square_diff(nums1: Vec, nums2: Vec, k1: i32, k2: i32) -> i64 { + Self::min_sum_square_diff(nums1, nums2, k1, k2) + } +} + +#[cfg(test)] +mod tests { + #[test] + fn test_solution() { + super::super::tests::run::(); + } +}