diff --git a/src/problem/mod.rs b/src/problem/mod.rs index dc58e39..122d838 100644 --- a/src/problem/mod.rs +++ b/src/problem/mod.rs @@ -180,4 +180,5 @@ mod p33_search_in_rotated_sorted_array; mod p34_find_first_and_last_position_of_element_in_sorted_array; mod p153_find_minimum_in_rotated_sorted_array; mod p215_kth_largest_element_in_an_array; -mod p502_ipo; \ No newline at end of file +mod p502_ipo; +mod p373_find_k_pairs_with_smallest_sums; \ No newline at end of file diff --git a/src/problem/p373_find_k_pairs_with_smallest_sums.rs b/src/problem/p373_find_k_pairs_with_smallest_sums.rs new file mode 100644 index 0000000..3909a88 --- /dev/null +++ b/src/problem/p373_find_k_pairs_with_smallest_sums.rs @@ -0,0 +1,86 @@ +/** + * [373] Find K Pairs with Smallest Sums + */ +pub struct Solution {} + + +// submission codes start here +use std::cmp::Ordering; + +#[derive(Eq, PartialEq)] +struct Pair<'a> { + nums1: &'a Vec, + nums2: &'a Vec, + x: usize, + y: usize, +} + +impl Pair<'_> { + fn new<'a>(nums1: &'a Vec, nums2: &'a Vec, x: usize, y: usize) -> Pair<'a> { + Pair { + nums1, + nums2, + x, + y, + } + } +} + +impl Ord for Pair<'_> { + fn cmp(&self, other: &Self) -> Ordering { + (other.nums1[other.x] + other.nums2[other.y]).cmp( + &(self.nums1[self.x] + self.nums2[self.y])) + } +} + +impl PartialOrd for Pair<'_> { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Solution { + pub fn k_smallest_pairs(nums1: Vec, nums2: Vec, k: i32) -> Vec> { + use std::collections::BinaryHeap; + + let (m, n) = (nums1.len(), nums2.len()); + let k = k as usize; + + let mut heap = BinaryHeap::new(); + + for i in 0..m.min(k) { + heap.push(Pair::new(&nums1, &nums2, i, 0)); + } + + let mut count = 0; + let mut result = Vec::with_capacity(k); + + while count < k && !heap.is_empty() { + let pair = heap.pop().unwrap(); + let (x, y) = (pair.x, pair.y); + result.push((nums1[x], nums2[y])); + + if y + 1 < n { + heap.push(Pair::new(&nums1, &nums2, x, y + 1)); + } + + count += 1; + } + + + result.iter().map(|(x, y)| vec![*x, *y]).collect() + } +} + +// submission codes end + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_373() { + assert_eq!(vec![vec![1, 2], vec![1, 4], vec![1, 6]], + Solution::k_smallest_pairs(vec![1, 7, 11], vec![2, 4, 6], 3)); + } +}