Add problem 2040: Kth Smallest Product of Two Sorted Arrays

src/lib.rs

@@ -1520,6 +1520,7 @@ pub mod problem_2035_partition_array_into_two_arrays_to_minimize_sum_difference;
 pub mod problem_2037_minimum_number_of_moves_to_seat_everyone;
 pub mod problem_2038_remove_colored_pieces_if_both_neighbors_are_the_same_color;
 pub mod problem_2039_the_time_when_the_network_becomes_idle;
+pub mod problem_2040_kth_smallest_product_of_two_sorted_arrays;
 pub mod problem_2042_check_if_numbers_are_ascending_in_a_sentence;
 pub mod problem_2043_simple_bank_system;
 pub mod problem_2044_count_number_of_maximum_bitwise_or_subsets;

src/problem_2040_kth_smallest_product_of_two_sorted_arrays/binary_search.rs

@@ -0,0 +1,161 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+use std::mem;
+
+impl Solution {
+    fn check(
+        (left_negative, left_zero, left_positive): (&[i32], &[i32], &[i32]),
+        (right_negative, right_zero, right_positive): (&[i32], &[i32], &[i32]),
+        k: u64,
+        middle: i64,
+    ) -> bool {
+        let mut count = 0;
+
+        if middle < 0 {
+            let mut iter = right_positive.iter();
+
+            'outer: for &lhs in left_negative {
+                loop {
+                    if let Some(&first) = iter.as_slice().first() {
+                        if i64::from(lhs) * i64::from(first) <= middle {
+                            count += iter.len() as u64;
+
+                            break;
+                        }
+
+                        iter.next();
+                    } else {
+                        break 'outer;
+                    }
+                }
+            }
+
+            iter = right_negative.iter();
+
+            'outer: for &lhs in left_positive.iter().rev() {
+                loop {
+                    if let Some(&last) = iter.as_slice().last() {
+                        if i64::from(lhs) * i64::from(last) <= middle {
+                            count += iter.len() as u64;
+
+                            break;
+                        }
+
+                        iter.next_back();
+                    } else {
+                        break 'outer;
+                    }
+                }
+            }
+        } else {
+            count += left_negative.len() as u64 * (right_zero.len() + right_positive.len()) as u64;
+            count += left_zero.len() as u64 * (right_negative.len() + right_zero.len() + right_positive.len()) as u64;
+            count += left_positive.len() as u64 * (right_negative.len() + right_zero.len()) as u64;
+
+            let mut iter = right_negative.iter();
+
+            'outer: for &lhs in left_negative.iter().rev() {
+                loop {
+                    if let Some(&first) = iter.as_slice().first() {
+                        if i64::from(lhs) * i64::from(first) <= middle {
+                            count += iter.len() as u64;
+
+                            break;
+                        }
+
+                        iter.next();
+                    } else {
+                        break 'outer;
+                    }
+                }
+            }
+
+            iter = right_positive.iter();
+
+            'outer: for &lhs in left_positive {
+                loop {
+                    if let Some(&last) = iter.as_slice().last() {
+                        if i64::from(lhs) * i64::from(last) <= middle {
+                            count += iter.len() as u64;
+
+                            break;
+                        }
+
+                        iter.next_back();
+                    } else {
+                        break 'outer;
+                    }
+                }
+            }
+        }
+
+        count < k
+    }
+
+    fn min_max(mut a: i64, mut b: i64, mut c: i64, mut d: i64) -> (i64, i64) {
+        fn sort_2(a: &mut i64, b: &mut i64) {
+            if b < a {
+                mem::swap(a, b);
+            }
+        }
+
+        sort_2(&mut a, &mut b);
+        sort_2(&mut c, &mut d);
+
+        (a.min(c), b.max(d))
+    }
+
+    fn split_at_partition_point(nums: &[i32], f: impl FnMut(&i32) -> bool) -> (&[i32], &[i32]) {
+        let i = nums.partition_point(f);
+
+        nums.split_at(i)
+    }
+
+    fn split_3(nums: &[i32]) -> (&[i32], &[i32], &[i32]) {
+        let (negative, rest) = Self::split_at_partition_point(nums, |&x| x < 0);
+        let (zero, positive) = Self::split_at_partition_point(rest, |&x| x == 0);
+
+        (negative, zero, positive)
+    }
+
+    pub fn kth_smallest_product(nums1: Vec<i32>, nums2: Vec<i32>, k: i64) -> i64 {
+        let k = k as u64;
+        let min_1 = i64::from(*nums1.first().unwrap());
+        let max_1 = i64::from(*nums1.last().unwrap());
+        let min_2 = i64::from(*nums2.first().unwrap());
+        let max_2 = i64::from(*nums2.last().unwrap());
+        let (mut left, mut right) = Self::min_max(min_1 * min_2, min_1 * max_2, max_1 * min_2, max_1 * max_2);
+        let split_1 = Self::split_3(&nums1);
+        let split_2 = Self::split_3(&nums2);
+
+        while left < right {
+            let middle = (left + right).div_euclid(2);
+
+            if Self::check(split_1, split_2, k, middle) {
+                left = middle + 1;
+            } else {
+                right = middle;
+            }
+        }
+
+        left
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn kth_smallest_product(nums1: Vec<i32>, nums2: Vec<i32>, k: i64) -> i64 {
+        Self::kth_smallest_product(nums1, nums2, k)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}

src/problem_2040_kth_smallest_product_of_two_sorted_arrays/mod.rs

@@ -0,0 +1,30 @@
+pub mod binary_search;
+
+pub trait Solution {
+    fn kth_smallest_product(nums1: Vec<i32>, nums2: Vec<i32>, k: i64) -> i64;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            ((&[2, 5] as &[_], &[3, 4] as &[_], 2), 8),
+            ((&[-4, -2, 0, 3], &[2, 4], 6), 0),
+            ((&[-2, -1, 0, 1, 2], &[-3, -1, 2, 4, 5], 3), -6),
+            (
+                (
+                    &[-10, -9, -8, -5, -3, -2, 1, 2, 4, 8],
+                    &[-9, -8, -8, -4, -4, -3, -1, 0, 4],
+                    73,
+                ),
+                32,
+            ),
+        ];
+
+        for ((nums1, nums2, k), expected) in test_cases {
+            assert_eq!(S::kth_smallest_product(nums1.to_vec(), nums2.to_vec(), k), expected);
+        }
+    }
+}