Find K Pairs with Smallest Sums

题目地址:
https://leetcode.com/problems/find-k-pairs-with-smallest-sums/description/

题目:
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:
Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

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:
Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

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:
Given nums1 = [1,2], nums2 = [3],  k = 3 

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]

解题思路:
这道题就是用优先队列的办法来做。首先先把num1里面的所有元素和num2里面的第一个元素来配对并且将index设为0。

代码:

class Pair{
    int[] pair;
    int idx; // current index to nums2    long sum;
    Pair(int idx, int n1, int n2){
        this.idx = idx;
        pair = new int[]{n1, n2};
        sum = (long) n1 + (long) n2;
    }
}

public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
    List<int[]> ret = new ArrayList<>();
    if (nums1==null || nums2==null || nums1.length ==0 || nums2.length ==0) return ret;
    int len1 = nums1.length, len2=nums2.length;

    PriorityQueue<Pair> q = new PriorityQueue<>(k, new Comparator() {
        @Override        public int compare(Object o1, Object o2) {
            return (int) (((Pair)o1).sum - ((Pair)o2).sum);
        }
    });
    for (int i=0; i<nums1.length && i<k ; i++) { // only need first k number in nums1 to start        q.offer( new Pair(0, nums1[i],nums2[0]) );
    }
    for (int i=1; i<=k && !q.isEmpty(); i++) { // get the first k sums        Pair p = q.poll();
        ret.add( p.pair );
        if (p.idx < len2 -1 ) { // get to next value in nums2            int next = p.idx+1;
            q.offer( new Pair(next, p.pair[0], nums2[next]) );
        }
    }
    return ret;
}








Comments

Popular Posts