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;
    }
};

Link

LeetCode-373 Find K Pairs with Smallest Sums