simple Java O(KlogK) solution with explanation


  • 109
    B

    Basic idea: Use min_heap to keep track on next minimum pair sum, and we only need to maintain K possible candidates in the data structure.

    Some observations: For every numbers in nums1, its best partner(yields min sum) always strats from nums2[0] since arrays are all sorted; And for a specific number in nums1, its next candidate sould be [this specific number] + nums2[current_associated_index + 1], unless out of boundary;)

    Here is a simple example demonstrate how this algorithm works.

    image

    The run time complexity is O(kLogk) since que.size <= k and we do at most k loop.

    public class Solution {
        public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
            PriorityQueue<int[]> que = new PriorityQueue<>((a,b)->a[0]+a[1]-b[0]-b[1]);
            List<int[]> res = new ArrayList<>();
            if(nums1.length==0 || nums2.length==0 || k==0) return res;
            for(int i=0; i<nums1.length && i<k; i++) que.offer(new int[]{nums1[i], nums2[0], 0});
            while(k-- > 0 && !que.isEmpty()){
                int[] cur = que.poll();
                res.add(new int[]{cur[0], cur[1]});
                if(cur[2] == nums2.length-1) continue;
                que.offer(new int[]{cur[0],nums2[cur[2]+1], cur[2]+1});
            }
            return res;
        }
    }
    

  • 0
    S

    @bbccyy1 best solution, so far. So easy to understand. Thanks!


  • 1
    G

    what is the name of this type of code "((a,b)->a[0]+a[1]-b[0]-b[1])". I am trying to learn this. It looks like a definition of comparator.


  • 6

  • 1
    C

    I believe the time complexity is not kLogK, but an amortized k*klogK solution. The best case would be klogk + (k-1)log(k-1) + (k-2)log(k-2) + ... + 1(In every iteration, the queue size only polls but not offer), the worst case is k*klogk.


  • 0
    I

    @chs5003 yep, it should be K*KlogK, because the loop iterates k times, and each time is KlogK


  • 1
    B

    @infini In JAVA, PriorityQueue is implemented as a min_heap through which one can achieve logN time to extract min and insert an element. It's actually a tree structure. More information about min heap.


  • 0
    C

    @bbccyy1 Okay, that makes sense.


  • 0
    I

    @bbccyy1 Totally true. Thanks for your reply!


  • 1

    Thx for sharing. If permuting them all in a matrix, you will find the problem is select kth smallest from num1.length sorted arrays. You can shorten your explanation. :-)


  • 0
    O

    Thanks for sharing. This is my C++ version.

    class Solution {
    private:
        struct Compare {
            bool operator()(tuple<int, int, int>& t1, tuple<int, int, int>& t2) {
                return ((get<0>(t1) + get<1>(t1)) > (get<0>(t2) + get<1>(t2)));
            }
        };
    
    public:
        vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
    
            vector<pair<int, int>> ans;
            
            if(nums1.size() == 0 || nums2.size() == 0 || k == 0) return ans;
            
            priority_queue<tuple<int, int, int>, vector<tuple<int, int, int> >, Compare> pq;
            
            for(int i = 0; i < min((int)nums1.size(), k); i++) pq.push(make_tuple(nums1[i], nums2[0], 0));
            
            while(k-- > 0 && !pq.empty()) {
                tuple<int, int, int> top = pq.top();
                ans.push_back(make_pair(get<0>(top), get<1>(top)));
                
                pq.pop();
                if(get<2>(top) + 1 < nums2.size())
                    pq.push(make_tuple(get<0>(top), nums2[get<2>(top)+1], get<2>(top)+1));
            }
            return ans;
        }
    };
    

  • 2

    Thanks for the figure. IT'S SOOO CLEAR.


  • 0

    @chs5003 It is O(kLogk). Look at the code:

    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        if (nums1.length == 0 || nums2.length == 0 || k == 0) return new ArrayList<>();
    
        List<int[]> result = new ArrayList<>();
        PriorityQueue<int[]> que = new PriorityQueue<>((a,b) -> (a[0] + a[1]) - (b[0] + b[1]));
        
        // runs min(k, nums1.length) times -> 1, log2, log2, log3 ... logk
        for(int i = 0; i < Math.min(nums1.length, k); i++)
            que.offer(new int[]{nums1[i], nums2[0], 0});
        
        // runs k times -> logk, logk, logk ...
        while(k-- > 0 && !que.isEmpty()){
            int[] cur = que.poll(); // pop the min
            result.add(new int[]{ cur[0], cur[1] }); // add the first 2 columns of min to result
            
            if (cur[2] == nums2.length-1) continue; // if (nums2Index == end of nums2) then continue
            
            int currNums1Val = cur[0];
            int nextNums2Index = cur[2] + 1;
            int nextNums2Value = nums2[nextNums2Index];
            
            que.offer(new int[]{ currNums1Val, nextNums2Value, nextNums2Index });
        }
        
        return result;
    }
    

  • 0
    A
    This post is deleted!

  • 0
    R

    Thank you so much


  • 0
    M

    nice solution


  • 1
    K

    How are we choosing to iterate through num1 rather than num2 in this step
    for(int i=0; i<nums1.length && i<k; i++) que.offer(new int[]{nums1[i], nums2[0], 0});

    Can you please explain this reasoning behind this?


  • 0
    H

    If you don't mind, can you please explain me the comparison written for PriorityQueue int he fallowing step: (a,b)->a[0]+a[1]-b[0]-b[1])


  • 0
    E

    It's kklog k. Just because the min heap has a capacity of k doesn't mean you are only iterating over k elements when adding to it. You are generating combinations of pairs from n1 and n2 and then adding each one one of them to the heap. Assuming worst case where n1.length == n2.length then it's kk which gives the time complexity of k^2. For each of those you add it to the heap which is log n therefore you have kk*log k.


  • 0
    M

    I like the idea of using the three-tuple as the entry for PrioirityQueue, with the last item in the tuple maintaining the index!


Log in to reply
 

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