Leetcode-373
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.
乍看之下似乎很简单，基于给定数组构建最小堆，删除k次堆顶即可
public class Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
List<List<Integer>> result = new ArrayList<>();
if (nums1.length == 0 || nums2.length == 0 || k == 0) {
return result;
}
PriorityQueue<Pair> pq = new PriorityQueue<>(
Comparator.comparingInt(i -> i.num1 + i.num2));
for (int i = 0; i < nums1.length; i++) {
for (int j = 0; j < nums2.length; j++) {
pq.add(new Pair(nums1[i], nums2[j]));
}
}
while (k-- > 0 && !pq.isEmpty()) {
Pair pair = pq.poll();
result.add(pairToList(pair));
}
return result;
}
private List<Integer> pairToList(Pair pair) {
return Arrays.asList(pair.num1, pair.num2);
}
private class Pair {
private int num1;
private int num2;
private Pair(int num1, int num2) {
this.num1 = num1;
this.num2 = num2;
}
}
}提交这段代码，居然直接AC，只能说Leetcode的OJ对超时的判定很宽松。
上述解法有两个问题：第一是时间复杂度是O(nums1.length*nums2.length)，如果数组很大，k很小，就会显得很低效；第二是这个解法根本没有利用到数组有序的这个条件，经验告诉我们，如果已知条件没有充分利用，这个解法很可能不是最优的。
为了利用有序这个条件，我们注意到：
设序列的当前元素为
(nums1[i1], nums2[i2])，由于数组有序，序列的下一个元素一定是(nums1[i1+1], nums2[i2])或者(nums1[i1], nums2[i2+1])
根据以上命题，我们优化一下上面的答案，首先给辅助类Pair添加数组下标字段
private class Pair {
private int i1; //num1在数组1的下标
private int i2; //num2在数组2的下标
private int num1;
private int num2;
private Pair(int i1,int i2, int num1, int num2) {
this.i1 = i1;
this.i2 = i2;
this.num1 = num1;
this.num2 = num2;
}
}把序列的第一个元素放入堆中后，进行k次循环，每次删除堆顶，得到当前序列元素，然后把下一个可能的序列元素放入堆中。
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
List<List<Integer>> result = new ArrayList<>();
if (nums1.length == 0 || nums2.length == 0 || k == 0) {
return result;
}
PriorityQueue<Pair> pq = new PriorityQueue<>(
Comparator.comparingInt(i -> i.num1 + i.num2));
pq.add(new Pair(0, 0, nums1[0], nums2[0]));
for (int i = 0; i < nums1.length * nums2.length && i < k; i++) {
Pair pair = pq.poll();
result.add(pairToList(pair));
if (pair.i2 + 1 < nums2.length) {
pq.add(new Pair(pair.i1, pair.i2 + 1, nums1[pair.i1], nums2[pair.i2 + 1]));
}
if (pair.i1 + 1 < nums1.length) {
pq.add(new Pair(pair.i1 + 1, pair.i2, nums1[pair.i1 + 1], nums2[pair.i2]));
}
}
return result;
}似乎优化得不错，利用到了数组有序的条件，且外层循环的时间复杂度为O(k)，循环内堆添加元素和删除堆顶操作的时间复杂度均为O(logk)，总的时间复杂度为O(klogk)。提交代码，发现WA，有重复的元素。
举例说明上面代码的问题，简单起见用(0,0)表示(nums1[0],nums2[0])，以此类推：
- 初始堆：(0,0)
- 第一次循环，删除堆顶，添加元素(0,1) (1,0)，得到堆(0,1) (1,0)
- 第二次循环，假定堆顶是(0,1)，添加元素(0,2) (1,1)，得到堆 (1,0) (0,2) (1,1)
- 第三次循环，假定堆顶是(1,0)，添加元素(1,1) (2,0)，得到堆 (0,2) (1,1) (1,1) (2,0)
在步骤4中，(1,1)被添加了2次，这就是为什么上述解法可能会导致重复元素。
下面的代码参考了讨论区其他人的解法，用到了多路归并排序的思想，可以解决元素重复添加的问题。
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
List<List<Integer>> result = new ArrayList<>();
if (nums1.length == 0 || nums2.length == 0 || k == 0) {
return result;
}
PriorityQueue<Pair> pq = new PriorityQueue<>(
Comparator.comparingInt(i -> i.num1 + i.num2));
//第一轮，把元素(0,0) (1,0) (2,0) ...(k-1,0)放入堆中
//相当于构建了k个子序列
for (int i = 0; i < k && i < nums1.length; i++) {
pq.add(new Pair(i, 0, nums1[i], nums2[0]));
}
for (int i = 0; i < nums1.length * nums2.length && i < k; i++) {
//删除堆顶操作相当于查看k个子序列中，每个子序列的第1个元素，再删除其中的最小值
//最小值记为(i,j)
Pair pair = pq.poll();
result.add(pairToList(pair));
//i对应的子序列第1个元素被删除后，原第2个元素就变成了第1个元素
if (pair.i2 + 1 < nums2.length) {
pq.add(new Pair(pair.i1, pair.i2 + 1, nums1[pair.i1], nums2[pair.i2 + 1]));
}
}
return result;
}