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