Find K Pairs with Smallest Sums LeetCode solution in Rust

1
impl Solution {
2
  pub fn k_smallest_pairs(nums1: Vec<i32>, nums2: Vec<i32>, k: i32) -> Vec<Vec<i32>> {
3
    use std::cmp::Reverse;
4
    use std::collections::BinaryHeap;
5

6
    let mut heap: BinaryHeap<(Reverse<i32>, usize, usize)> = (0..nums1.len())
7
      .map(|i| (Reverse(nums1[i] + nums2[0]), i, 0))
8
      .collect();
9
    let mut ret = Vec::with_capacity(k as usize);
10

11
    while let Some((_, i, j)) = heap.pop() {
12
      ret.push(vec![nums1[i], nums2[j]]);
13
      if ret.len() == k as usize {
14
        break;
15
      }
16

17
      if j + 1 < nums2.len() {
18
        heap.push((Reverse(nums1[i] + nums2[j + 1]), i, j + 1));
19
      }
20
    }
21
    ret
22
  }
23
}