373. Find K Pairs with Smallest Sums

Question

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 (u****1****,v****1****),(u****2****,v****2****) ...(u****k****,v****k****) 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]

---

Code

public class Solution {
    public class Tuple {
        public int index1;
        public int index2;
        public int value1;
        public int value2;
        public Tuple(int index1, int index2, int value1, int value2) {
            this.index1 = index1;
            this.index2 = index2;
            this.value1 = value1;
            this.value2 = value2;
        }
    }
    
    public List kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List result = new ArrayList<>();
        if (k == 0 || nums1.length == 0 || nums2.length == 0) return result;
        
        PriorityQueue pq = new PriorityQueue<>(new Comparator() {
            public int compare(Tuple t1, Tuple t2) {
                return (t1.value1 + t1.value2) - (t2.value1 + t2.value2);
            }    
        });
        for (int i = 0; i < nums1.length; i++) {
            Tuple tuple = new Tuple(i, 0, nums1[i], nums2[0]);
            pq.offer(tuple);
        }
        while (pq.size() > 0 && result.size() < k) {
            Tuple tuple = pq.poll();
            
            int[] t = new int[2];
            t[0] = tuple.value1;
            t[1] = tuple.value2;
            result.add(t);
            
            if (tuple.index2 + 1 < nums2.length) {
                tuple.index2++;
                tuple.value2 = nums2[tuple.index2];
                pq.offer(tuple);
            }
        }
        
        return result;
    }
}

---

Solution

使用PriorityQueue解决。

自定义一个二元组Tuple，其中含有nums1中某元素的下标index1和值value1以及nums2中某元素的下标index2和值value2。

定义一个PriorityQueue，排序是按照Tuple中value1和value2的和排序。

初始化时，将nums1中所有元素与nums2中第一个元素组成Tuple，加入队列中。

循环，弹出队列中第一个元素，将元素的value1和value2加入结果集中。并将元素的index2 + 1，value2更新后，重新加入队列中。循环直至结果集大小为k或队列为空。