373 Find K Pairs with Smallest Sums

1. Question

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:
1
Given nums1 = [1,7,11], nums2 = [2,4,6], k = 3
2
3
Return: [1,2],[1,4],[1,6]
4
5
The first 3 pairs are returned from the sequence:
6
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
Copied!
Example 2:
1
Given nums1 = [1,1,2], nums2 = [1,2,3], k = 2
2
3
Return: [1,1],[1,1]
4
5
The first 2 pairs are returned from the sequence:
6
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
Copied!
Example 3:
1
Given nums1 = [1,2], nums2 = [3], k = 3
2
3
Return: [1,3],[2,3]
4
5
All possible pairs are returned from the sequence:
6
[1,3],[2,3]
Copied!

2. Implementation

(1) Heap
1
class Solution {
2
public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
3
List<int[]> res = new ArrayList<>();
4
5
if (nums1 == null || nums1.length == 0 || nums2 == null || nums2.length == 0) {
6
return res;
7
}
8
9
PriorityQueue<PairSum> minHeap = new PriorityQueue<>();
10
11
int m = nums1.length, n = nums2.length;
12
13
for (int j = 0; j < n; j++) {
14
minHeap.add(new PairSum(0, j, nums1[0] + nums2[j]));
15
}
16
17
for (int i = 0; i < Math.min(k, m * n); i++) {
18
PairSum curSum = minHeap.remove();
19
int index1 = curSum.index1;
20
int index2 = curSum.index2;
21
22
res.add(new int[] {nums1[index1], nums2[index2]});
23
24
if (index1 < m - 1) {
25
minHeap.add(new PairSum(index1 + 1, index2, nums1[index1 + 1] + nums2[index2]));
26
}
27
}
28
return res;
29
}
30
31
class PairSum implements Comparable<PairSum> {
32
int index1;
33
int index2;
34
int sum;
35
36
public PairSum(int index1, int index2, int sum) {
37
this.index1 = index1;
38
this.index2 = index2;
39
this.sum = sum;
40
}
41
42
public int compareTo(PairSum that) {
43
return this.sum - that.sum;
44
}
45
}
46
}
Copied!

3. Time & Space Complexity

Heap: 时间复杂度O(k * k *logk), 空间复杂度O(k)