【資料結構】演算法 Find K Pairs with Smallest Sums 查詢和最小的K對數字
阿新 • • 發佈:2021-10-11
目錄
Find K Pairs with Smallest Sums 查詢和最小的K對數字
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.
Return the k pairs (u1, v1), (u2, v2), ..., (uk, vk) with the smallest sums.
nums1 和nums2都是有序的,找最小sum的pair
輸入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
輸出: [1,2],[1,4],[1,6]
解釋: 返回序列中的前 3 對數:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
思路
可以通過大頂堆,併為每一個pair構造陣列,1,2存num,3存sum
並且將堆大小保持在k。此時堆內所有陣列節點都是需要的結果。
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) { PriorityQueue<List<Integer>> pq = new PriorityQueue<List<Integer>>((o1, o2)->o2.get(2)-o1.get(2)); List<List<Integer>> ans = new ArrayList<>(); for (int i = 0; i < nums1.length; i++) { for (int j = 0; j < nums2.length; j++) { int sum = nums1[i] +nums2[j]; if(pq.size()<k||sum <pq.peek().get(2)){ pq.offer(new ArrayList<>(Arrays.asList(nums1[i],nums2[j],sum))); if(pq.size()>k){ pq.poll(); } } else{ break; } } } while (pq.size()>k){ pq.poll(); } for (List<Integer> integers : pq) { List<Integer> item = new ArrayList<>(); item.add(integers.get(0)); item.add(integers.get(1)); ans.add(item); } return ans; }
Tag
heap