python:
class Solution:
def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
queue = []
# 入堆
def push(i,j):
if i<len(nums1) and j<len(nums2):
heapq.heappush(queue,[nums1[i]+nums2[j],i,j])
# 将两个数组的第一个元素加入
for i in range(len(nums1)):
push(i,0)
pairs = []
while queue and len(pairs) < k:
_,i,j = heapq.heappop(queue)
pairs.append([nums1[i],nums2[j]])
push(i,j+1)
return pairs
Java:
class Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
// 定义优先队列
PriorityQueue<int[]> queue = new PriorityQueue<>((o1,o2)->(nums1[o1[0]]+nums2[o1[1]]) - (nums1[o2[0]]+nums2[o2[1]]));
List<List<Integer>> res = new LinkedList<>();
// 任意数组为空,返回空
if(nums1.length == 0 || nums2.length == 0){
return res;
}
//将每个数组的第一项加入小顶堆中
// nums1中的每个元素和nums2中的第一个元素【坐标为0】
for(int i = 0;i<Math.min(nums1.length,k);i++){
queue.add(new int[]{i,0});
}
// 循环,只要堆不为空 并且 k 个数字
while(k>0 && !queue.isEmpty()){
// 弹出堆顶元素
int[] pair = queue.poll();
// [[1,2],[1,4],[1,6]]
List<Integer> item = new ArrayList<>();
item.add(nums1[pair[0]]);
item.add(nums2[pair[1]]);
// 将当前元素的下一个元素加入堆
if(pair[1]<nums2.length-1){
queue.add(new int[]{pair[0],pair[1]+1});
}
res.add(item);
k--;
}
return res;
}
}