題目敘述

題目難度：Medium
題目描述：題目給定兩個陣列 nums1 以及 nums2，兩者都以 non-decreasing order 排序，並且給定整數 k，題目要求你從兩個陣列中個取出一個整數，定義成 pair (u, v) 請找出 k 具有最小總和的 pair u1, v1), (u2, v2), ..., (uk, vk)

解法

一開始的想法

一開始的想法比較暴力一點，一共三步：

兩個陣列都各自迭代找出所有 pair 組合
定義 minHeap 和 comparator 來比較pair之間誰的總和比較小
pop k 次

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<vector<int>> resultPairs;
        //define heap
        auto comparator = [](const pair<int, int> &leftPair, const pair<int, int> &rightPair){
            int sumLeft = leftPair.first +  leftPair.second;
            int sumRight = rightPair.first + rightPair.second;
            return sumLeft > sumRight;
        };
        priority_queue<pair<int,int>, vector<pair<int,int>>, decltype(comparator)> minHeap(comparator);
        
        // form the pairs
        for(int i=0; i<nums1.size();i++){
            for(int j=0; j<nums2.size();j++){
                // vecPairs.push_back({nums1[i], nums2[j]});
                minHeap.push({nums1[i],nums2[j]});
            }
        }

        //pop k times
        while(!minHeap.empty() && k>0){
            resultPairs.push_back({minHeap.top().first, minHeap.top().second});
            minHeap.pop();
            k--;
        }

        return resultPairs;
    }
};

這是原先寫的方法，但這種方法絕對爆炸，因為會找出 nums1.size() * nums2.size() 個pairs 記憶體一定炸開，會Memory Limit Exceeded

我的解法

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<vector<int>> resultPairs;

        //define heap
        auto comparator = [&](const pair<int, int> &leftPair, const pair<int, int> &rightPair){
            long long sumLeft = (long long)(nums1[leftPair.first] +  nums2[leftPair.second]);
            long long sumRight = (long long)(nums1[rightPair.first] + nums2[rightPair.second]);
            return sumLeft > sumRight;
        };
        priority_queue<pair<int,int>, vector<pair<int,int>>, decltype(comparator)> minHeap(comparator);
        
        // form the pairs
        int limit = min((int)nums1.size(), k);
        for(int i=0; i<limit; i++){
            minHeap.push({i,0});
        }

        while(!minHeap.empty() && k-->0){
            auto [i, j] = minHeap.top();
            minHeap.pop();
            resultPairs.push_back({nums1[i],nums2[j]});
            if (j + 1 < (int)nums2.size()) {
                minHeap.push({i, j + 1});
            }
        }
        return resultPairs;
    }
};