[LeetCode] (medium) 373. Find K Pairs with Smallest Sums

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

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]

---

逻辑很简单，但是实现有一些关于stl的trivial的细节

优先级队列priority_queue<>的第三个模板参数，必须是仿函数

class Solution {
public:
    class myComparison{
    public:
        bool operator()(pair<int, int> &P1, pair<int, int> &P2) const {
            return ((P1.first+P1.second) <= (P2.first+P2.second));
        }
    };
    
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        priority_queue<pair<int, int>, vector<pair<int, int>>, myComparison> PQ;
        
        for(int i = 0; i < nums1.size(); ++i){
            for(int j = 0; j < nums2.size(); ++j){
                pair<int, int> tem(nums1[i], nums2[j]);
                
                if(PQ.size() < k) PQ.push(tem);
                else{
                    pair<int, int> peek(PQ.top());
                    if(myComparison()(tem, peek)){
                        PQ.pop();
                        PQ.push(tem);
                    }else{
                        break;
                    }
                }
            }
        }
        
        vector<pair<int, int>> result;
        while(!PQ.empty()){
            result.push_back(PQ.top());
            PQ.pop();
        }
        
        return result;
    }
    
};

看到了别人的做法，他的遍历方法很酷，充分利用了两个vector的有序性，也防止了重复访问的情况。

直观上可以理解为一个三角形，每一层的二元组的两个数的和相同，向下延伸一层的时候就相当于把当前层向左下方平移一格，再在新层的最右侧添加一个结点。

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int,int>> result;
        if (nums1.empty() || nums2.empty() || k <= 0)
            return result;
        auto comp = [&nums1, &nums2](pair<int, int> a, pair<int, int> b) {
            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];};
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(comp)> min_heap(comp);
        min_heap.emplace(0, 0);
        while(k-- > 0 && min_heap.size())
        {
            auto idx_pair = min_heap.top(); min_heap.pop();
            result.emplace_back(nums1[idx_pair.first], nums2[idx_pair.second]);
            if (idx_pair.first + 1 < nums1.size())
                min_heap.emplace(idx_pair.first + 1, idx_pair.second);
            if (idx_pair.first == 0 && idx_pair.second + 1 < nums2.size())
                min_heap.emplace(idx_pair.first, idx_pair.second + 1);
        }
        return result;
    }
};