題目要求
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.
兩個單調遞增的整數數組,現分別從數組1和數組2中取一個數字構成數對,求找到k個和最小的數對。
思路
這題採用最大堆作爲輔助的數據結構能夠完美的解決我們的問題。觀察數組我們可以看到,從nums1中任意取一個數字,其和nums2中的數字組成的最小數對一定是<nums1[k], nums2[0]>
,同理,我們可以知道,<nums1[k], nums2[t+1]>
的值一定比nums1[k], nums2[t]
大。因此在優先隊列中,我們存儲所有的nums1中數字所能夠構成的最小數對。每從堆中取走一個數對<nums1[k], nums2[t]>
,就插入<nums1[k], nums2[t+1]>
,從而確保堆中的數對都可以從小到大遍歷到。
public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
List<int[]> result = new ArrayList<int[]>();
if(nums1.length == 0 || nums2.length == 0 || k == 0) return result;
PriorityQueue<int[]> heap = new PriorityQueue<int[]>(new Comparator<int[]>(){
@Override
public int compare(int[] o1, int[] o2) {
return o1[0] + o1[1] - o2[0] - o2[1];
}});
for(int i = 0 ; i<nums1.length ; i++){
heap.offer(new int[]{nums1[i], nums2[0], 0});
}
while(k-- != 0 && !heap.isEmpty()) {
int[] min = heap.poll();
result.add(new int[]{min[0], min[1]});
if(min[2] == nums2.length) continue;
heap.offer(new int[]{min[0], nums2[min[2]+1], min[2] + 1});
}
return result;
}