You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.
Define a pair (u,v) which consists of one element from the first array and one element from the second array.
Find the k pairs(u1,v1),(u2,v2) ...(uk,vk)with the smallest sums.
Example 1:
Given nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Return: [1,2],[1,4],[1,6]
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]
Example 2:
Given nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Return: [1,1],[1,1]
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]
Example 3:
Given nums1 = [1,2], nums2 = [3], k = 3
Return: [1,3],[2,3]
All possible pairs are returned from the sequence:
[1,3],[2,3]
2. Implementation
(1) Heap
class Solution {
public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
List<int[]> res = new ArrayList<>();
if (nums1 == null || nums1.length == 0 || nums2 == null || nums2.length == 0) {
return res;
}
PriorityQueue<PairSum> minHeap = new PriorityQueue<>();
int m = nums1.length, n = nums2.length;
for (int j = 0; j < n; j++) {
minHeap.add(new PairSum(0, j, nums1[0] + nums2[j]));
}
for (int i = 0; i < Math.min(k, m * n); i++) {
PairSum curSum = minHeap.remove();
int index1 = curSum.index1;
int index2 = curSum.index2;
res.add(new int[] {nums1[index1], nums2[index2]});
if (index1 < m - 1) {
minHeap.add(new PairSum(index1 + 1, index2, nums1[index1 + 1] + nums2[index2]));
}
}
return res;
}
class PairSum implements Comparable<PairSum> {
int index1;
int index2;
int sum;
public PairSum(int index1, int index2, int sum) {
this.index1 = index1;
this.index2 = index2;
this.sum = sum;
}
public int compareTo(PairSum that) {
return this.sum - that.sum;
}
}
}