Description
给定两个已排序的数组\(nums1[]\)和\(nums2[]\)，和一个常数\(k\)。设元组\(p(u, v)\)为从\(nums1[]\)拿取\(u\)，且从\(nums2[]\)拿取\(v\)
返回\(k\)个元组\(p(u_i, v_j)\)使得它们的和最小
Constraints
- \(nums1[]\)和\(nums2[]\)长度不超过\(10^5\)
- \(nums1[i]\)和\(nums2[i]\)都在
int范围内 - \(k\)不超过\(10^4\)
Solution
假设你拿了\(u_i\)和\(v_j\)，那么可以确定：
- \(u_i + v_j \le u_i + v_{j+1}\)
- \(u_i + v_j \le u_{i+1} + v_j\)
但是你不知道\(u_i + v_j\)对比\(u_{i+1} + v_0\)或者\(u_0 + v_{j+1}\)的大小关系
所以初始化就是把所有\(p(u_0, v_j)\)和\(p(u_i, v_0)\)都放到堆里，拿走一个最优解后各下标偏移1位
注意不要重复放取，因此加上一个vis表
Code
class Solution {
public:
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
using Tup = tuple<int, int>;
auto cmp = [&](auto &a, auto &b) {
auto &[a1, a2] = a;
auto &[b1, b2] = b;
return nums1[a1] + nums2[a2] > nums1[b1] + nums2[b2];
};
priority_queue<Tup, vector<Tup>, decltype(cmp)> pq(cmp);
unordered_map<size_t, unordered_map<size_t, bool>> vis;
for(size_t i = 0; i < nums1.size(); ++i) {
vis[i][0] = true;
pq.emplace(make_tuple(i, 0));
}
for(size_t j = 1; j < nums2.size(); ++j) {
vis[0][j] = true;
pq.emplace(make_tuple(0, j));
}
vector<vector<int>> ans;
while(ans.size() < k && !pq.empty()) {
auto [i, j] = pq.top();
pq.pop();
vector<int> best {nums1[i], nums2[j]};
ans.emplace_back(move(best));
const int d[] = {1, 0};
for(auto t = 0; t < 2; ++t) {
auto ni = i + d[t];
auto nj = j + d[t^1];
if(ni >= nums1.size() || nj >= nums2.size()) continue;
if(vis[ni][nj]) continue;
vis[ni][nj] = true;
pq.emplace(make_tuple(ni, nj));
}
}
return ans;
}
};