正文
[LeetCode] 373. 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.
Find the k pairs (u1,v1),(u2,v2) ...(uk,vk) with the smallest sums.
Example 1:
Given nums1 = [1,7,11], nums2 = [2,4,6], k = 3 Return: [1,2],[1,4],[1,6] The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
Example 2:
Given nums1 = [1,1,2], nums2 = [1,2,3], k = 2 Return: [1,1],[1,1] The first 2 pairs are returned from the sequence:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
Example 3:
Given nums1 = [1,2], nums2 = [3], k = 3 Return: [1,3],[2,3] All possible pairs are returned from the sequence:
[1,3],[2,3]
Credits:
Special thanks to @elmirap and @StefanPochmann for adding this problem and creating all test cases.
给定2个以升序排列的整数数组和1个整数k，从2个数组中各拿出一个数字组成一对，找出k对和最小的数字组合。
解法1：最简单的想法就是暴力brute force解法，但效率肯定不高。
解法2: 最小堆。把所有的点对加入到最小堆，然后输出前k个。但没有利用到“两个数组都有序”这个条件，就算数组无序，也可以利用这个方法。要利用有序这个条件，可以借助mergesort的思路，pair的第一个元素至多包含了nums1数组的前k个元素，k以后的可以不用考虑。所以，这形成了k个list，每一个list都包含了nums2的元素。每一次取所有list中的最小值，然后该list下一个元素入队。
Java：
public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
List<int[]> r = new ArrayList<>();
if(nums1.length == 0 || nums2.length == 0) return r;
int size = Math.min(nums1.length, k);
int[] index = new int[size];
PriorityQueue<int[]> queue = new PriorityQueue<>(new Comparator<int[]>(){
@Override
public int compare(int[] o1, int[] o2) {
Integer s1 = o1[0] + o1[1];
Integer s2 = o2[0] + o2[1];
return s1.compareTo(s2);
}
});
for(int i = 0; i < size; i++){
queue.add(new int[]{nums1[i], nums2[0], i});
}
int count = 0;
while(!queue.isEmpty()){
int[] pair = queue.poll();
r.add(new int[]{pair[0], pair[1]});
int id = pair[2];
if(++index[id] < nums2.length)
queue.add(new int[]{nums1[id], nums2[index[id]], id});
count++;
if(count == k)
break;
}
return r;
}
Python:
from heapq import heappush, heappop class Solution(object):
def kSmallestPairs(self, nums1, nums2, k):
"""
:type nums1: List[int]
:type nums2: List[int]
:type k: int
:rtype: List[List[int]]
"""
pairs = []
if len(nums1) > len(nums2):
tmp = self.kSmallestPairs(nums2, nums1, k)
for pair in tmp:
pairs.append([pair[1], pair[0]])
return pairs min_heap = []
def push(i, j):
if i < len(nums1) and j < len(nums2):
heappush(min_heap, [nums1[i] + nums2[j], i, j]) push(0, 0)
while min_heap and len(pairs) < k:
_, i, j = heappop(min_heap)
pairs.append([nums1[i], nums2[j]])
push(i, j + 1)
if j == 0:
push(i + 1, 0) # at most queue min(n, m) space
return pairs
C++:
// Time: O(k * log(min(n, m, k))), where n is the size of num1, and m is the size of num2.
// Space: O(min(n, m, k)) class Solution {
public:
vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
vector<pair<int, int>> pairs;
if (nums1.size() > nums2.size()) {
vector<pair<int, int>> tmp = kSmallestPairs(nums2, nums1, k);
for (const auto& pair : tmp) {
pairs.emplace_back(pair.second, pair.first);
}
return pairs;
} using P = pair<int, pair<int, int>>;
priority_queue<P, vector<P>, greater<P>> q;
auto push = [&nums1, &nums2, &q](int i, int j) {
if (i < nums1.size() && j < nums2.size()) {
q.emplace(nums1[i] + nums2[j], make_pair(i, j));
}
}; push(0, 0);
while (!q.empty() && pairs.size() < k) {
auto tmp = q.top(); q.pop();
int i, j;
tie(i, j) = tmp.second;
pairs.emplace_back(nums1[i], nums2[j]);
push(i, j + 1);
if (j == 0) {
push(i + 1, 0); // at most queue min(m, n) space.
}
}
return pairs;
}
};
