An improved solution with min-heap, c++ beats 87.77%


  • 0
    E

    This problem can be changed into Kth Smallest Element in a Sorted Matrix
    As the T(n) of the min-heap solution for 'kth sorted matrix' is mainly related to the min-heap size, so we can use the smaller of (row and column) to create heap.

    typedef vector<int>::iterator ittype;
    typedef pair<ittype,ittype> ptype;
    struct compare{ //Store iterator pairs,because I don't want to store four number:(i,nums1[i],j,num2[j])
        bool operator()(ptype&p1,ptype&p2){
            return (*p1.first+*p1.second)>(*p2.first+*p2.second);
        }
    }; 
    class Solution {
    public:
        vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
            priority_queue<int,vector<ptype>,compare>pq;
            vector<pair<int,int> >res;
            if(nums1.size()==0||nums2.size()==0)return res;
            ittype it1,it_end1,it2,it_end2;
            bool f=0;
            //it1 points to the shorter array, it2 points to the longer, and the heap size is the same as the shorter.
            if(nums1.size()>nums2.size()){it1=nums2.begin();it_end1=nums2.end();it2=nums1.begin();it_end2=nums1.end();f=1;}
            else{it1=nums1.begin();it_end1=nums1.end();it2=nums2.begin();it_end2=nums2.end();}
            for(;it1!=it_end1;it1++)pq.push(ptype(it1,it2));
            while(!pq.empty()&&k>0){
                k--;
                ptype p=pq.top();
                //Because the TEST sets make nums1 first and nums2 second,if I submit (nums2,nums1), it can't pass, so I reverse such.
                if(!f)res.push_back(pair<int,int>(*p.first,*p.second));
                else res.push_back(pair<int,int>(*p.second,*p.first));
                pq.pop();
                if(p.second+1==it_end2)continue;
                pq.push(ptype(p.first,p.second+1));
            }
            return res;
        }
    };
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.