Heap-based Java Solution, similar to MERGING K SORTED LIST


  • 1
    F

    The idea lying behind this problem is the same as merging K sorted lists:
    List 1: a1+b1, a1+b2, a1+b3,a1+b4,a1+b5,...........
    List 2: a2+b1, a2+b2, a2+b3,a2+b4,a2+b5,...........
    List 3: a3+b1, a3+b2, a3+b3,a3+b4,a3+b5,...........
    ......
    So, first we need to add the first element in every list to the heap. Then, we keep popping out the top of the heap and adding the next element in the list which the currently popped out element comes from, until we get k smallest pairs or all the pairs are used up. The time complexity is: O(max(length1, k) * log(length1)), where length1 is the length of the first ascending array.

    public class Solution {
        public List<int[]> kSmallestPairs(final int[] nums1, final int[] nums2, int k) {
            ArrayList<int[]> topKInts = new ArrayList<>();
            PriorityQueue<int[]> pq;
            
            if(nums1.length==0 || nums2.length==0) {
                return topKInts;
            }
            
            pq = new PriorityQueue(nums1.length, new Comparator<int[]>() {
                @Override
                public int compare(int[] indices1, int[] indices2) {
                    if(nums1[ indices1[ 0 ] ]+nums2[ indices1[ 1 ] ] < nums1[ indices2[ 0 ] ]+nums2[ indices2[ 1 ] ]) {
                        return -1;
                    } else if(nums1[ indices1[ 0 ] ]+nums2[ indices1[ 1 ] ] > nums1[ indices2[ 0 ] ]+nums2[ indices2[ 1 ] ]) {
                        return 1;
                    } else {
                        return 0;
                    }
                }
            });
            
            for(int cnt = 0; cnt < nums1.length; cnt++) {
                pq.offer(new int[]{cnt, 0});
            }
            
            while(!pq.isEmpty() && topKInts.size()<k) {
                int[] indices = pq.poll();
                topKInts.add(new int[]{nums1[ indices[ 0 ] ], nums2[ indices[ 1 ] ]});
                indices[ 1 ]++;
                if(indices[ 1 ] < nums2.length) {
                    pq.offer(indices);
                }
            }
            
            return topKInts;
        }
    }
    

Log in to reply
 

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