82 lines
1.8 KiB
Rust
82 lines
1.8 KiB
Rust
/**
|
|
* [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<i32>,
|
|
nums2: &'a Vec<i32>,
|
|
x: usize,
|
|
y: usize,
|
|
}
|
|
|
|
impl Pair<'_> {
|
|
fn new<'a>(nums1: &'a Vec<i32>, nums2: &'a Vec<i32>, 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<Ordering> {
|
|
Some(self.cmp(other))
|
|
}
|
|
}
|
|
|
|
impl Solution {
|
|
pub fn k_smallest_pairs(nums1: Vec<i32>, nums2: Vec<i32>, k: i32) -> Vec<Vec<i32>> {
|
|
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)
|
|
);
|
|
}
|
|
}
|