Leetcode-373
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;
}