373. Find K Pairs with Smallest Sums**

373. Find K Pairs with Smallest Sums**

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/description/

题目描述

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u,v) which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1),(u2,v2) ...(uk,vk) with the smallest sums.

Example 1:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]] 
Explanation: The first 3 pairs are returned from the sequence: 
             [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [1,1],[1,1]
Explanation: The first 2 pairs are returned from the sequence: 
             [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [1,3],[2,3]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

解题思路

这道题最直观的思路是使用一个最大堆, 保存 pair<int, vector<int>>, 第一个值表示 sum = u + v, 第二个值表示 (nums[i], nums[j]). 首先在堆中放入 k 个值, 对于新来的值, 如果它比堆中的最大元素还小, 就插入到堆中.

C++ 实现 1

priority_queue 默认是最大堆.

lass Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {       
        vector<vector<int>> res;
        int count = 0;  // count 来记录当前堆中是否插入了 k 个元素.
        priority_queue<pair<int, vector<int>>> Q;
        for (int i = 0; i < nums1.size(); ++i) {
            for (int j = 0; j < nums2.size(); ++j) {
                if (count < k) {
                    Q.push(make_pair(nums1[i] + nums2[j], vector<int>{nums1[i], nums2[j]}));
                    count ++;
                } else { // 如果新来的元素小于堆的 top, 那么插入新元素, 弹出旧元素.
                    if (nums1[i] + nums2[j] < Q.top().first) {
                        Q.pop();
                        Q.push(make_pair(nums1[i] + nums2[j], vector<int>{nums1[i], nums2[j]}));
                    } else break; // 这个 else 分支去掉也没有问题
                }
            }
        }
        while (!Q.empty()) {
            auto p = Q.top();
            res.push_back(p.second);
            Q.pop();
        }
        return res;
    }
};