题目描述

给定两个以升序排列的整形数组 nums1 和 nums2, 以及一个整数 k。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2。

找到和最小的 k 对数字 (u1, v1), (u2, v2), …, (uk, vk)。

样例

输入: 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]

输入: 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]

输入: nums1 = [1,2], nums2 = [3], k = 3 
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]

算法1

(二分) $O(\log S \cdot n \cdot \log m + k)$

1. 我们通过二分找到第 k 个数对对应的和是多少。这一步查找可以通过枚举每一个 nums1 的元素，然后用 upper_bound 求出有多少个数对满足小于等于 mid。假设这个和是 l。
2. 然后我们暴力加入所有小于 l 的数对到答案中，最多有 k 个。然后用类似的方法将等于 l 的数对加入到答案中（对于每个 nums1 的元素，通过 upper_bound 和 lower_bound 的差值来确定有多少个等于 l 的数对）。

时间复杂度

• 二分套 upper_bound 的时间复杂度为 $O(\log S \cdot n \cdot \log m)$，之后暴力统计答案的时间复杂度为 $O(k + n \log m)$。
• 故总时间复杂度为 $O(\log S \cdot n \cdot \log m + k)$，其中 $S$ 为最大的两个数之和与最小的两个数之和的差值，$n$ 和 $m$ 可以替换。
• 此方法适用于 $k$ 较大，$n$ 和 $m$ 的中有一个比较小的情况。

空间复杂度

• 只需存储答案，故空间复杂度为 $O(k)$。

C++ 代码

class Solution {
public:
    int calc(int x, const vector<int>& nums1, const vector<int>& nums2) {
        int n = nums1.size(), m = nums2.size(), cnt = 0;
        for (int i = 0; i < n; i++)
            cnt += upper_bound(nums2.begin(), nums2.end(), x - nums1[i]) - nums2.begin();

        return cnt;
    }

    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        int n = nums1.size(), m = nums2.size();
        vector<vector<int>> ans;

        if (n == 0 || m == 0)
            return ans;

        int l = nums1.front() + nums2.front(), r = nums1.back() + nums2.back();
        while (l < r) {
            int mid = (l + r) >> 1;
            if (calc(mid, nums1, nums2) >= k) {
                r = mid;
            } else {
                l = mid + 1;
            }
        }

        for (int i = 0; i < n; i++) {
            int j = 0;
            while (j < m && nums1[i] + nums2[j] < l) {
                ans.push_back({nums1[i], nums2[j]});
                j++;
            }
        }


        for (int i = 0; i < n; i++) {
            int x = l - nums1[i];
            int cnt = upper_bound(nums2.begin(), nums2.end(), x) - lower_bound(nums2.begin(), nums2.end(), x);
            if (ans.size() + cnt >= k)
                cnt = k - ans.size();


            for (int j = 0; j < cnt; j++)
                ans.push_back({nums1[i], x});

            if (ans.size() == k)
                break;
        }

        return ans;
    }
};

算法2

(堆) $O(k \log k)$

1. 我们可以将这个问题看成 m 个有序链表合并，求出前 k 个元素的问题。
2. 所以我们可以将所有 (0, y) 数对放入堆中。每次出堆时加入答案。假设出堆的是 (i, y)，若 i < n - 1，则 (i + 1, y) 加入队列。直到出堆了 k 次为止。

时间复杂度

• 堆中最多有个 $k$ 个元素，每次操作的时间复杂度为 $O(\log k)$，共 $k$ 次，故时间复杂度为 $O(k \log k)$。
• 此方法适用于 $k$ 较小，$n$ 和 $m$ 较大。

空间复杂度

• 需要额外 $O(k)$ 的空间维护堆，答案也需要 $O(k)$ 的空间。

C++ 代码

struct A {
    int x, y, i;
    A(int x_, int y_, int i_): x(x_), y(y_), i(i_) {}
    bool operator < (const A &other) const {
        return x + y > other.x + other.y;
    }
};

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        int n = nums1.size(), m = nums2.size();
        vector<vector<int>> ans;

        if (n == 0 || m == 0)
            return ans;

        priority_queue<A> heap;

        for (int i = 0; i < min(k, m); i++)
            heap.push(A(nums1[0], nums2[i], 0));

        while (!heap.empty()) {
            if (ans.size() == k)
                return ans;

            auto s = heap.top();
            heap.pop();
            ans.push_back({s.x, s.y});
            if (s.i < n - 1)
                heap.push(A(nums1[s.i + 1], s.y, s.i + 1));
        }
        return ans;
    }
};