Find K Pairs with Smallest Sums

Problem

You are given two integer arrays and sorted in non-decreasing order and an integer .

Define a pair which consists of one element from the first array and one element from the second array.

Return the pairs with the smallest sums.

Example

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]
Explanation: The first 3 pairs are returned from the sequence: [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]
Explanation: The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [[1,3],[2,3]]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

Constraints

  • and both are sorted in non-decreasing order

Submit your solution at here

Solution

Intuition

At first I tried to use 2 pointers, but failed to do so. Then I thought about priority queue

Approach

  • Use a priority queue to manage a list of candidates
  • Init the queue with element
  • Pop the queue to get the lowest element. Let that element is we will push element and element to the queue and continue to the next iteration

Complexity

  • Time complexity:
  • Space complexity:

Code

use std::collections::BinaryHeap;
use std::collections::HashSet;
use std::cmp::Reverse;
impl Solution {
    pub fn k_smallest_pairs(nums1: Vec<i32>, nums2: Vec<i32>, k: i32) -> Vec<Vec<i32>> {
        let n = nums1.len();
        let m = nums2.len();
        let mut q = BinaryHeap::new();
        let mut vis = HashSet::new();
        let i = 0;
        let j = 0;
        let x = Reverse(nums1[i]+nums2[j]);
        q.push((x,i,j));
        while let Some((Reverse(x), i, j)) = q.pop() {
            if vis.contains(&(i,j)) {
                continue;
            }
            vis.insert((i,j));
            if vis.len() >= k as usize {
                break;
            }
            if i < n-1 {
                let x = Reverse(nums1[i+1]+nums2[j]);
                let next = (x, i+1, j);
                q.push(next)
            }
            if j < m-1 {
                let x = Reverse(nums1[i]+nums2[j+1]);
                let next = (x, i, j+1);
                q.push(next)
            }
        }
        vis.into_iter()
            .map(|(i,j)| vec![nums1[i], nums2[j]])
            .collect()
    }
}