Add problem 2333: Minimum Sum of Squared Difference

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;

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<i32>, nums2: Vec<i32>, k1: i32, k2: i32) -> i64;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        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);
+        }
+    }
+}

src/problem_2333_minimum_sum_of_squared_difference/quick_select.rs

@@ -0,0 +1,158 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+struct Random<const N: usize>(usize);
+
+impl<const N: usize> Random<N> {
+    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<i32>, nums2: Vec<i32>, 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::<Vec<_>>();
+        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<i32>, nums2: Vec<i32>, 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::<super::Solution>();
+    }
+}