题目

给定两个以升序排列的整数数组 nums1 和 nums2 , 以及一个整数 k 。定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。请找到和最小的 k 个数对 (u1,v1),  (u2,v2)  ...  (uk,vk) 。示例 1:
输入: 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]示例 2:
输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数：[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]示例 3:
输入: nums1 = [1,2], nums2 = [3], k = 3 
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]提示:
1 <= nums1.length, nums2.length <= 104
-109 <= nums1[i], nums2[i] <= 109
nums1, nums2 均为升序排列
1 <= k <= 1000

题解

C++

class Solution {
public:vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {auto cmp = [&nums1, &nums2](const pair<int, int> & a, const pair<int, int> & b) {return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];};int m = nums1.size();int n = nums2.size();vector<vector<int>> ans;   priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> pq(cmp);for (int i = 0; i < min(k, m); i++) {pq.emplace(i, 0);}while (k-- > 0 && !pq.empty()) {auto [x, y] = pq.top(); pq.pop();ans.emplace_back(initializer_list<int>{nums1[x], nums2[y]});if (y + 1 < n) {pq.emplace(x, y + 1);}}return ans;}
};

java

class Solution {public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {PriorityQueue<int[]> pq = new PriorityQueue<>(k, (o1, o2)->{return nums1[o1[0]] + nums2[o1[1]] - nums1[o2[0]] - nums2[o2[1]];});List<List<Integer>> ans = new ArrayList<>();int m = nums1.length;int n = nums2.length;for (int i = 0; i < Math.min(m, k); i++) {pq.offer(new int[]{i,0});}while (k-- > 0 && !pq.isEmpty()) {int[] idxPair = pq.poll();List<Integer> list = new ArrayList<>();list.add(nums1[idxPair[0]]);list.add(nums2[idxPair[1]]);ans.add(list);if (idxPair[1] + 1 < n) {pq.offer(new int[]{idxPair[0], idxPair[1] + 1});}}return ans;}
}

C++

class Solution {
public:vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {int m = nums1.size();int n = nums2.size();auto count = [&](int target){long long cnt = 0;int start = 0;int end = n - 1;while (start < m && end >= 0) {if (nums1[start] + nums2[end] > target) {end--;} else {cnt += end + 1;start++;}}return cnt;};/*二分查找第 k 小的数对和的大小*/int left = nums1[0] + nums2[0];int right = nums1.back() + nums2.back();int pairSum = right;while (left <= right) {int mid = left + ((right - left) >> 1);if (count(mid) < k) {left = mid + 1;} else {pairSum = mid;right = mid - 1;}}vector<vector<int>> ans;int pos = n - 1;/*找到小于目标值 pairSum 的数对*/for (int i = 0; i < m; i++) {while (pos >= 0 && nums1[i] + nums2[pos] >= pairSum) {pos--;}for (int j = 0; j <= pos && k > 0; j++, k--) {ans.push_back({nums1[i], nums2[j]});}}/*找到等于目标值 pairSum 的数对*/pos = n - 1;for (int i = 0; i < m && k > 0; i++) {while (pos >= 0 && nums1[i] + nums2[pos] > pairSum) {pos--;}for (int j = i; k > 0 && j >= 0 && nums1[j] + nums2[pos] == pairSum; j--, k--) {ans.push_back({nums1[i], nums2[pos]});}}return ans;}
};

Java

class Solution {public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {int m = nums1.length;int n = nums2.length;/*二分查找第 k 小的数对和的大小*/int left = nums1[0] + nums2[0];int right = nums1[m - 1] + nums2[n - 1];int pairSum = right;while (left <= right) {int mid = left + ((right - left) >> 1);long cnt = 0;int start = 0;int end = n - 1;while (start < m && end >= 0) {if (nums1[start] + nums2[end] > mid) {end--;} else {cnt += end + 1;start++;}}if (cnt < k) {left = mid + 1;} else {pairSum = mid;right = mid - 1;}}List<List<Integer>> ans = new ArrayList<>();int pos = n - 1;/*找到小于目标值 pairSum 的数对*/for (int i = 0; i < m; i++) {while (pos >= 0 && nums1[i] + nums2[pos] >= pairSum) {pos--;}for (int j = 0; j <= pos && k > 0; j++, k--) {List<Integer> list = new ArrayList<>();list.add(nums1[i]);list.add(nums2[j]);ans.add(list);}}/*找到等于目标值 pairSum 的数对*/pos = n - 1;for (int i = 0; i < m && k > 0; i++) {while (pos >= 0 && nums1[i] + nums2[pos] > pairSum) {pos--;}for (int j = i; k > 0 && j >= 0 && nums1[j] + nums2[pos] == pairSum; j--, k--) {List<Integer> list = new ArrayList<>();list.add(nums1[j]);list.add(nums2[pos]);ans.add(list);}}return ans;}
}