373. Find K Pairs with Smallest Sums (java,优先队列)

时间:2022-01-09 06:13:44

题目:

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<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
if(nums1 == null || nums2 == null || k <= 0)
return null;
if(nums1.length == 0 || nums2.length == 0)
return new ArrayList<int[]>();
if(nums1.length * nums2.length < k)
k = nums1.length * nums2.length;
//the priorityqueue 优先队列
PriorityQueue<Pair> queue = new PriorityQueue<Pair>();
//result
List<int[]> list = new ArrayList<int[]>();
//首先将nums1同nums2[0]的所有组合入队列
for(int i = 0; i < nums1.length; ++i)
queue.add(new Pair(i,0,nums1[i],nums2[0]));
//获取k个最小值
for(int n = 0; n < k; ++n){
Pair tmp = queue.poll();
int[] ele = tmp.pair; //数值对
list.add(ele); //加入结果集
if(tmp.j == nums2.length - 1)
continue;
queue.add(new Pair(tmp.i,tmp.j + 1,nums1[tmp.i],nums2[tmp.j + 1]));
}
return list;
}
}
//定义pair的数据结构
class Pair implements Comparable<Pair>{
int[] pair = new int[2]; //定义数组对
int i; //第一个元素在nums1中的下标
int j; //第二个元素在nums2中的下标
public Pair(int i,int j,int e1,int e2){
pair[0] = e1;
pair[1] = e2;
this.i = i;
this.j = j;
} @Override
public int compareTo(Pair that){ //升序排列
return (pair[0] + pair[1]) - (that.pair[0] + that.pair[1]);
}
}